EnzymeAD / EnzymeAD/Reactant.jl

Using `ComponentArrays` fails to compile a JAX function

Open
#1,344 2 comments 0 reactions 0 assignees View on GitHub
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.