EnzymeAD / EnzymeAD/Reactant.jl
Using `ComponentArrays` fails to compile a JAX function
- Dominant language
- Julia
- Stars
- 370
- Forks
- 74
- Avg merge
- 18h 47m
- Merged PRs (30d)
- 30
Description
First of all, this is a fascinating feature. Thanks to all the team who made it possible.
I am interested in using `ComponentArrays` since the elements can be accessed through their names. I try to use a `ComponentVector`, convert it to a reactant array and compile a user-defined jax function.
```julia
using PythonCall, Reactant, ComponentArrays
jax = pyimport("jax")
np = pyimport("numpy")
namespace = pydict()
namespace["jax"] = jax
jax_func_str = """
def f(x, t):
a, w, p = x
return a * jax.numpy.sin(w * t + p)
"""
pyexec(jax_func_str, namespace)
f = namespace["f"]
t = [1.0]
p_comparr = ComponentVector(a=1.0, w=2.0, p=0.0)
t_rarr = Reactant.to_rarray(t)
p_comparr_rarr = Reactant.to_rarray(p_comparr)
f_text = jax.jit(f).lower(np.array(p_comparr), np.array(t)).as_text()
f_comp = @compile f(p_comparr_rarr, t_rarr) # Fails due to scalar indexing
@compile Reactant.Ops.hlo_call(pyconvert(String, f_text), p_comparr_rarr, t_rarr) # all inputs should be reactant arrays
```
This line `f_comp = @compile f(p_comparr_rarr, t_rarr)` fails by emitting the following error
```julia
ERROR: Python: Julia: Scalar indexing is disallowed.
Invocation of getindex(::TracedRArray, ::Vararg{Int, N}) resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore should be avoided.
If you want to allow scalar iteration, use `allowscalar` or `@allowscalar`
to enable scalar iteration globally or for the operations in question.
```
Compiling via `hlo_call` complains that not all the inputs are of reactant array type.
```julia
>julia @compile Reactant.Ops.hlo_call(pyconvert(String, f_text), p_comparr_rarr, t_rarr) # all inputs should be reactant arrays
ERROR: AssertionError: hlo_call: all inputs to hlo_call should be reactant arrays or numbers
```
I suspect that this error is due to this warning when I convert `p_comparr` to a reactant array
```julia
julia> p_comparr_rarr = Reactant.to_rarray(p_comparr)
┌ Warning: `Adapt.parent_type` is not implemented for Vector{Float64}. Assuming Vector{Float64} isn't a wrapped array.
└ @ Reactant ~/.julia/packages/Reactant/lLXhp/src/Reactant.jl:43
ComponentVector{Float64, ConcretePJRTArray{Float64, 1, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}, Tuple{Axis{(a = 1, w = 2, p = 3)}}}(a = 1.0, w = 2.0, p = 0.0)
```
The same example works without any issues if I use Julia arrays
```julia
using PythonCall, Reactant, ComponentArrays
jax = pyimport("jax")
np = pyimport("numpy")
namespace = pydict()
namespace["jax"] = jax
jax_func_str = """
def f(x, t):
a, w, p = x
return a * jax.numpy.sin(w * t + p)
"""
pyexec(jax_func_str, namespace)
f = namespace["f"]
t = [1.0]
p_arr = [1.0, 2.0, 0.0]
t_rarr = Reactant.to_rarray(t)
p_rarr = Reactant.to_rarray(p_arr)
f_text = jax.jit(f).lower(np.array(p_arr), np.array(t)).as_text()
f_comp = @compile f(p_rarr, t_rarr)
@compile Reactant.Ops.hlo_call(pyconvert(String, f_text), p_rarr, t_rarr)
```
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.