Calling Python function JIT compiled with JAX from Julia without overhead

[sefffal]

It could be very useful to call JIT compiled JAX functions from Python libraries from Julia, and in turn, to be able for Julia libraries to call user-provided JAX functions.

This is possible today via PythonCall, e.g.:

    using PythonCall
    jax = pyimport("jax")
    numpy = pyimport("numpy")
    
    # Define a simple test function
    func_str = """
    def f(x):
        return jax.numpy.sin(x) + jax.numpy.cos(x)
    """
    
    # Create namespace and define function
    namespace = pydict()
    namespace["jax"] = jax
    pyexec(func_str, namespace)
    py_func = namespace["f"]
    
    # Create test input
    x = numpy.array([1.0], dtype=numpy.float32)
    
    # Get lowered representation
    lowered = jax.jit(py_func).lower(x)
    println("HLO text:")
    println(lowered.as_text())
    
    # Also test compilation and execution through Python
    compiled = lowered.compile()
    @time result = compiled(x)
    println("\nTest execution result:", result)
    println("\nTest execution result:", sum(result))

The problem of course it that we have overhead from python, about 7us worth per call, because starting the JAX function still goes through Python (compiled(x)...).

It seems like it should be possible to call the JIT compiled function directly. At worse case, it should be possible to export the compiled JAX function and then run it without involving Python.

Here is a relevant GitHub discussion: Calling pre-compiled JAX code from C++ · jax-ml/jax · Discussion #22184 · GitHub

Has anyone tried this, or have suggestions on how to proceed? I guess we would need to get the PJRT C library built and into Yggdrasil?

[wsmoses]

Might I recommend GitHub - EnzymeAD/Reactant.jl

cc @avikpal and @mofeing

there’s docs in lux (which we should move to reactant proper) for exporting julia into jax, we can make docs for doing the same in reverse. Exporting Lux Models to Jax (via EnzymeJAX & Reactant) | Lux.jl Docs

[sefffal]

Hi @wsmoses , thanks for this!

I started digging into the Reactant code.

Through PythonCall I can jit compile a function and get its “executable”.

After that, it looks like the reactant function codegen_xla_call is the next thing to use, right?

[wsmoses]

kind of, honestly I think the answer here is that we should just add an Ops.hlo_call that takes a string of a stablehlo module like we have in Enzyme-Jax (Enzyme-JAX/test/testffi.py at main · EnzymeAD/Enzyme-JAX · GitHub).

Then with that you could even make a jax call like this or something

function jax_call(py_func, args...)
    lowered = jax.jit(py_func).lower(args...)
    Ops.hlo_call(lowered.as_text(), args...)
end

x = Reactant.ConcreteRArray(ones(10))
jlfn = @compile jax_call(py_func, x)
jlfn(x)

If you’re interested in helping add this let me know and I can help you with the setup!

[sefffal]

I’ve spent some time reading up more on Reactant. Very cool!

Once a traced function is compiled with Reactant, is it cheap to call into that function? (Like a regular Julia function, without the cost imposed by starting a JAX kernel from Python). Ie the entire Julia code doesn’t have to be traced by Reactant to get good performance, right?
It looks like the answer is yes, which is great!

Similarly, if I use Reactant for one part of a computation, can one use enzyme to take gradients through a larger Julia program that happens to contain code compiled through reactant?

Thanks!

[wsmoses]

It should be super cheap to call into the function yeah. And correct, you compile/trace once and can re use it as much as you want (aka you won’t retrace it).

So currently the setup assumes that Reactant will be on the outside of an Enzyme autodiff (if the autodiff uses reactant types). We may be able to support the other direction, but you’re going to get significantly better perf by doing reactant on the outside anyways, so I’m curious on your use case here.

[sefffal]

Thanks, just trying to make get a clear picture of how it’s intended to work!

My main usecase is providing a Julia MCMC library to Python users. Ideally:

• we expose the library using JuliaCall eg like PySR
• they provide a JAX JIT traced Python function
• we run and autodiff their function using Reactant and Enzyme, ideally on multiple threads

I’m also interested to use some Python Gaussian process libraries written in JAX within a larger Julia program, which might not yet be amenable to fully tracing with Reactant (but maybe I’m wrong).

Sounds like almost everything is already in place for this, which is great! If I want to help creating a reactant op for general XLA code, where should I start?

[wsmoses]

so @Pangoraw just added support for hlo_call in this PR (Add Ops.hlo_call(::String, args...) by Pangoraw · Pull Request #358 · EnzymeAD/Reactant.jl · GitHub).

Probably the next thing to do in your case, is to add a Reactant extension for pythoncall of a TracedRArray and gets the stablehlo out of jax and instead does a corresponding Reactant.Ops.hlo_call

Want to give it a go? We’d be happy to help!

x/ref Trace over Python · Issue #354 · EnzymeAD/Reactant.jl · GitHub

[sefffal]

Okay this is pretty darn cool!

Simple proof of concept for tracing and JIT compiling a python function with JAX, and then running it in Julia and autodiffing it with Enzyme:

f, g = jax_from_julia_with_grad(
    """
    def f(x):
        return jax.numpy.sum(jax.numpy.sin(x) + jax.numpy.cos(x))
    """,
    ones(Float32, 100)
)
r = Reactant.to_rarray(ones(Float32, 100))

julia> @time f(r)
  0.000034 seconds (4 allocations: 80 bytes)
(ConcreteRNumber{Float32}(138.17734f0),)

julia> @time g(r)
  0.000043 seconds (4 allocations: 96 bytes)
(ConcreteRArray{Float32, 1}(Float32[-0.30116868, -0.30116868, -0.30116868, -0.30116868, -0.30116868, -0.30116868, -0.30116868, -0.30116868, -0.30116868, -0.30116868  …  -0.30116868, -0.30116868, -0.30116868, -0.30116868, -0.30116868, -0.30116868, -0.30116868, -0.30116868, -0.30116868, -0.30116868]),)

Code below:

using Reactant
using Reactant: Ops
using Enzyme
using PythonCall
function jax_from_julia_with_grad(func_str, example_inputs)

    # Create namespace and define function
    namespace = pydict()
    namespace["jax"] = jax
    pyexec(func_str, namespace)
    py_func = namespace["f"]
    
    # Create test input
    x = numpy.array(example_inputs, dtype=numpy.float32)
    
    # Get lowered representation
    lowered = jax.jit(py_func).lower(x)

    _primal = @compile Ops.hlo_call(
        pyconvert(String, lowered.as_text()),
        Reactant.to_rarray(example_inputs)    
    )

    function _grad(i)
        function inner(inp)
            Ops.hlo_call(
                pyconvert(String, lowered.as_text()),
                inp
            )
        end
        return Enzyme.gradient(Reverse, inner, i)
    end
    __grad = @compile _grad(
        Reactant.to_rarray(example_inputs)    
    )


    return (args...)->_primal(nothing,args...), __grad
end

[sefffal]

The above works (shockingly!) well as a way to call a given Python/JAX function without overhead.

Looking at how this could be integrated more smoothly into PythonCall + Reactant though, and it’s not clear that PythonCall has the necessary extension points.

The public interface to extend PythonCall is through adding new argument conversion rules. Here, we would need something that operates on an entire function call + any arguments that are TracedRArrays–something like:

using Reactant
using Reactant: Ops
using Enzyme
using PythonCall
jax = pyimport("jax")
numpy = pyimport("numpy")

function PythonCall.pycall(f::Py, args::Reactant.TracedRArray...; kwargs...)
    lowered = jax.jit(f).lower(args...)
    inputs = map(args) do arg
        numpy.array(size(arg), dtype=numpy.float32)
    end
    return Ops.hlo_call(
        pyconvert(String, lowered.as_text()),
        inputs...
    )
end

f = @compile jax.numpy.sum(
    Reactant.to_rarray(Float32[1, 2, 3]),
)

but this is reaching into internals (plus it stack overflows).

[wsmoses]

I mean worst case, we could add a python_call to Reactant without a definition unless Pythoncall is loaded.

Also for ease, mind moving this discussion to the relevant Reactant issue for visibility?

[wsmoses]

Following up here this now natively works in Reactant with optimization and differentiation between julia/python :

using Reactant
jax = pyimport("jax")

result = Reactant.@jit jax.numpy.sum(Reactant.to_rarray(Float32[1, 2, 3]))
@test typeof(result) == ConcreteRNumber{Float32}
@test result ≈ 6

from Reactant.jl/test/integration/python.jl at main · EnzymeAD/Reactant.jl · GitHub

[sefffal]

This is amazing, thanks @wsmoses and team!

I notice too that the overhead has decreased significantly since those experiments I posted above. They had a floor of about 30 us latency, and this example, when compiled, is more like 1.5 us. Amazing!

[ymardoukhi]

This looks amazing! I wonder if it would be possible to get the gradient of a compiled function as follows.

using Reactant
using Enzyme
jax = pyimport("jax")

t = Reactant.to_rarray([1.0])
dt = similar(t)
f = Reactant.@compile jax.numpy.sin(t)
Enzyme.autodiff(Reverse, f, Const, Duplicated(t, dt))

julia> Enzyme.autodiff(Reverse, f, Const, Duplicated(t, dt))
ERROR:
No augmented forward pass found for XLAExecuteSharded
 at context:   call void @XLAExecuteSharded(i64 %40, i32 noundef 1, [1 x i64]* nocapture noundef nonnull readonly %inpa.i.i, i64 %.fca.2.0.extract, [1 x i8]* nocapture noundef nonnull readonly %dona.i.i, i32 noundef 1, [1 x i64]* nocapture noundef nonnull writeonly %outa.i.i, i8* nocapture noundef nonnull writeonly %futa.i.i, [1 x i64]* nocapture noundef nonnull writeonly %futpa.i.i) #32, !dbg !128

Stacktrace:
 [1] macro expansion
   @ ~/.julia/packages/Reactant/lLXhp/src/xla/PJRT/LoadedExecutable.jl:195
 [2] execute_sharded
   @ ~/.julia/packages/Reactant/lLXhp/src/xla/PJRT/LoadedExecutable.jl:164
 [3] macro expansion
   @ ~/.julia/packages/Reactant/lLXhp/src/Compiler.jl:2731
 [4] Thunk
   @ ~/.julia/packages/Reactant/lLXhp/src/Compiler.jl:3159
 [5] Thunk
   @ ~/.julia/packages/Reactant/lLXhp/src/Compiler.jl:0

[avikpal]

@ymardoukhi

using Reactant
using Enzyme
jax = pyimport("jax")

t = Reactant.to_rarray([1.0])
dt = similar(t)
f = Reactant.@compile jax.numpy.sin(t)

Yes! Just compile the gradient function

julia> @code_hlo Enzyme.gradient(Reverse, jax.numpy.sum, t)
module @reactant_gradient attributes {mhlo.num_partitions = 1 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<1xf32>) -> tensor<1xf32> {
    %cst = stablehlo.constant dense<1.000000e+00> : tensor<1xf32>
    return %cst : tensor<1xf32>
  }
}

julia> @jit Enzyme.gradient(Reverse, jax.numpy.sum, t)
(ConcreteIFRTArray{Float32, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}, Nothing}(Float32[1.0]),)

[ymardoukhi]

Thanks! That was quite helpful. I was experimenting with ComponentArrays.jl, and I failed to compile my test function f (given below).

using PythonCall
using Enzyme
using Reactant
using 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]
t_rarr = Reactant.to_rarray(t)

################ ComponentArrays ###################
p_comparr = ComponentVector(a=1.0, w=2.0, p=0.0)
p_comparr_rarr = Reactant.to_rarray(p_comparr)

# This fails due to scalar indexing
f_comp = @compile f(p_comparr_rarr, t_rarr)

# This one fails with "ERROR: AssertionError: hlo_call: all inputs
# to hlo_call should be reactant arrays or numbers"
f_text = jax.jit(f).lower(np.array(p_comparr), np.array(t)).as_text()
@compile Reactant.Ops.hlo_call(pyconvert(String, f_text), p_comparr_rarr, t_rarr)

################ Arrays ##################
p_arr = [1.0, 2.0, 0.0]
p_rarr = Reactant.to_rarray(p_arr)

# This works
f_comp = @compile f(p_rarr, t_rarr)

# This works fine
f_text = jax.jit(f).lower(np.array(p_arr), np.array(t)).as_text()
@compile Reactant.Ops.hlo_call(pyconvert(String, f_text), p_rarr, t_rarr)

The error I get for f_comp = @compile f(p_comparr_rarr, t_rarr)

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.

And the error I get for @compile Reactant.Ops.hlo_call(pyconvert(String, f_text), p_comparr_rarr, t_rarr)

ERROR: AssertionError: hlo_call: all inputs to hlo_call should be reactant arrays or numbers

I am not sure what I am doing wrong here. I got the warning below when converting the ComponentVector to ConcretePJRTArray. Not sure if this is relevant.

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
AssertionError("Could not find registered platform with name: \"cuda\". Available platform names are: ")
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)

[mofeing]

For the first error, try wrapping f with a @allowscalar like this:

g(args...) = @allowscalar f(args...)

For the second error, do you mind opening an issue in Reactant.jl?

[ymardoukhi]

Thanks Sergio. I would like to avoid @allowscalar as I would like to target GPUs and TPUs. Nevertheless, using it throws the following error

julia> @compile g(p_comparr_rarr, t_rarr)
ERROR: Python: TypeError: Julia: MethodError: no method matching +(::Reactant.TracedRArray{Float64, 1}, ::Reactant.TracedRNumber{Float64})
For element-wise addition, use broadcasting with dot syntax: array .+ scalar
The function `+` exists, but no method is defined for this combination of argument types.

Closest candidates are:
  +(::Any, ::Any, !Matched::Any, !Matched::Any...)
   @ Base operators.jl:596
  +(!Matched::ChainRulesCore.NotImplemented, ::Any)
   @ ChainRulesCore ~/.julia/packages/ChainRulesCore/U6wNx/src/tangent_arithmetic.jl:24
  +(::Any, !Matched::ChainRulesCore.NotImplemented)
   @ ChainRulesCore ~/.julia/packages/ChainRulesCore/U6wNx/src/tangent_arithmetic.jl:25
  ...

This is certainly due to the argument of the sine function i.e. sin(w * t + p), which brings me to another question: is there a possibility to create a reactant number and not an array e.g. Reactant.to_rnumber(0.0)?

Also I opened up an issue on Reactant.jl.

[MilesCranmer]

Trying out this very cool example of JAX->Julia integration. I was wondering how you can get the latency down further? Might this be a regression or am I using things wrong?

e.g.,

julia> using Reactant: Reactant as Rx

julia> using PythonCall: pyimport

julia> using BenchmarkTools: @btime

julia> jax = pyimport("jax")
Python: <module 'jax' from '/private/var/folders/1h/xyppkvx52cl6w3_h8bw_gdqh0000gr/T/tmp.IbJRI7Nkqm/.CondaPkg/.pixi/envs/default/lib/python3.12/site-packages/jax/__init__.py'>

julia> jax_sum(x) = Rx.@jit jax.numpy.sum(x)
jax_sum (generic function with 1 method)

julia> @btime sum(x) setup=(x=randn(Float32,100))
  8.299 ns (0 allocations: 0 bytes)
-11.369346f0

julia> @btime jax_sum(x) setup=(x=Rx.to_rarray(randn(Float32,100)))
  42.593 ms (1615 allocations: 451.13 KiB)
Reactant.ConcretePJRTNumber{Float32, 1, Reactant.Sharding.ShardInfo{Reactant.Sharding.NoSharding, Nothing}}(19.730251f0)

This is on Julia 1.10. (1.11 had a precompilation error for ReactantPythonCallExt). Am using an M1 Pro in case that matters.

Edit: I tested with Rx.set_default_backend("cpu") but it didn’t help.

[Zentrik]

XLA itself does have quite a bit of overhead to call functions, but I think that should only be a couple microseconds (it’s been a while since I benchmarked it) and those allocations aren’t being counted by @btime.