# Calling Python function JIT compiled with JAX from Julia without overhead

**URL:** <https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552>\
**Category:** General Usage\
**Tags:** performance, interoperability\
**Created:** [December 6, 2024, 11:04pm UTC](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552 "2024-12-06T23:04:21Z")\
**Posts on this page:** 20\
**Page:** 1

<div class="post-metadata">

**Author:** ![sefffal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sefffal/32/23640_2.png) [@sefffal](https://discourse.julialang.org/u/sefffal)\
**Post date:** [December 6, 2024, 11:04pm UTC](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552/1 "2024-12-06T23:04:21Z")

</div>

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.:

```julia
    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](https://jax.readthedocs.io/en/latest/export/export.html) 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](https://github.com/jax-ml/jax/discussions/22184#discussioncomment-9909496)

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?

---

<div class="post-metadata">

**Author:** ![wsmoses](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/wsmoses/32/26497_2.png) [@wsmoses](https://discourse.julialang.org/u/wsmoses)\
**Post date:** [December 6, 2024, 11:24pm UTC](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552/2 "2024-12-06T23:24:07Z")

</div>

Might I recommend [GitHub - EnzymeAD/Reactant.jl](https://github.com/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](https://lux.csail.mit.edu/dev/manual/exporting_to_jax#Exporting-Lux-Models-to-Jax-(via-EnzymeJAX-and-Reactant))

---

<div class="post-metadata">

**Author:** ![sefffal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sefffal/32/23640_2.png) [@sefffal](https://discourse.julialang.org/u/sefffal)\
**Post date:** [December 7, 2024, 2:15am UTC](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552/3 "2024-12-07T02:15:05Z")

</div>

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?

---

<div class="post-metadata">

**Author:** ![wsmoses](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/wsmoses/32/26497_2.png) [@wsmoses](https://discourse.julialang.org/u/wsmoses)\
**Post date:** [December 7, 2024, 4:05am UTC](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552/4 "2024-12-07T04:05:26Z")

</div>

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](https://github.com/EnzymeAD/Enzyme-JAX/blob/main/test/testffi.py)).

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

```julia
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!

---

<div class="post-metadata">

**Author:** ![sefffal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sefffal/32/23640_2.png) [@sefffal](https://discourse.julialang.org/u/sefffal)\
**Post date:** [December 7, 2024, 4:11pm UTC](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552/5 "2024-12-07T16:11:45Z")

</div>

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!

---

<div class="post-metadata">

**Author:** ![wsmoses](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/wsmoses/32/26497_2.png) [@wsmoses](https://discourse.julialang.org/u/wsmoses)\
**Post date:** [December 10, 2024, 7:35am UTC](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552/6 "2024-12-10T07:35:06Z")

</div>

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.

---

<div class="post-metadata">

**Author:** ![sefffal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sefffal/32/23640_2.png) [@sefffal](https://discourse.julialang.org/u/sefffal)\
**Post date:** [December 10, 2024, 3:50pm UTC](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552/7 "2024-12-10T15:50:19Z")

</div>

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?

---

<div class="post-metadata">

**Author:** ![wsmoses](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/wsmoses/32/26497_2.png) [@wsmoses](https://discourse.julialang.org/u/wsmoses)\
**Post date:** [December 10, 2024, 8:16pm UTC](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552/8 "2024-12-10T20:16:31Z")

</div>

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](https://github.com/EnzymeAD/Reactant.jl/pull/358)).

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](https://github.com/EnzymeAD/Reactant.jl/issues/354)

---

<div class="post-metadata">

**Author:** ![sefffal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sefffal/32/23640_2.png) [@sefffal](https://discourse.julialang.org/u/sefffal)\
**Post date:** [December 10, 2024, 11:02pm UTC](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552/9 "2024-12-10T23:02:46Z")

</div>

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:

```julia
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:

```julia
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

```

---

<div class="post-metadata">

**Author:** ![sefffal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sefffal/32/23640_2.png) [@sefffal](https://discourse.julialang.org/u/sefffal)\
**Post date:** [December 10, 2024, 11:41pm UTC](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552/10 "2024-12-10T23:41:04Z")

</div>

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:

```julia
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).

---

<div class="post-metadata">

**Author:** ![wsmoses](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/wsmoses/32/26497_2.png) [@wsmoses](https://discourse.julialang.org/u/wsmoses)\
**Post date:** [December 11, 2024, 12:58am UTC](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552/11 "2024-12-11T00:58:21Z")

</div>

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?

---

<div class="post-metadata">

**Author:** ![wsmoses](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/wsmoses/32/26497_2.png) [@wsmoses](https://discourse.julialang.org/u/wsmoses)\
**Post date:** [December 30, 2024, 11:05pm UTC](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552/12 "2024-12-30T23:05:14Z")

</div>

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

```julia
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](https://github.com/EnzymeAD/Reactant.jl/blob/main/test/integration/python.jl)

---

<div class="post-metadata">

**Author:** ![sefffal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sefffal/32/23640_2.png) [@sefffal](https://discourse.julialang.org/u/sefffal)\
**Post date:** [December 31, 2024, 6:16pm UTC](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552/13 "2024-12-31T18:16:47Z")

</div>

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!

---

<div class="post-metadata">

**Author:** ![ymardoukhi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ymardoukhi/32/216992_2.png) [@ymardoukhi](https://discourse.julialang.org/u/ymardoukhi)\
**Post date:** [May 24, 2025, 1:52pm UTC](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552/14 "2025-05-24T13:52:42Z")

</div>

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

```julia
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
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

```

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [May 24, 2025, 2:21pm UTC](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552/15 "2025-05-24T14:21:59Z")

</div>

> [@ymardoukhi](#):
>
> ```julia
> 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
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]),)

```

---

<div class="post-metadata">

**Author:** ![ymardoukhi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ymardoukhi/32/216992_2.png) [@ymardoukhi](https://discourse.julialang.org/u/ymardoukhi)\
**Post date:** [May 24, 2025, 9:47pm UTC](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552/16 "2025-05-24T21:47:45Z")

</div>

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

```julia
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)`

```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.

```

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

```julia
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
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)

```

---

<div class="post-metadata">

**Author:** ![mofeing](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mofeing/32/47619_2.png) [@mofeing](https://discourse.julialang.org/u/mofeing)\
**Post date:** [May 27, 2025, 7:04pm UTC](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552/17 "2025-05-27T19:04:03Z")

</div>

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

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

```

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

---

<div class="post-metadata">

**Author:** ![ymardoukhi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/ymardoukhi/32/216992_2.png) [@ymardoukhi](https://discourse.julialang.org/u/ymardoukhi)\
**Post date:** [May 28, 2025, 9:15am UTC](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552/18 "2025-05-28T09:15:19Z")

</div>

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
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](https://github.com/EnzymeAD/Reactant.jl/issues/1344).

---

<div class="post-metadata">

**Author:** ![MilesCranmer](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/milescranmer/32/21070_2.png) [@MilesCranmer](https://discourse.julialang.org/u/MilesCranmer)\
**Post date:** [May 28, 2025, 5:25pm UTC](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552/19 "2025-05-28T17:25:21Z")

</div>

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-repl
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.

---

<div class="post-metadata">

**Author:** ![Zentrik](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/zentrik/32/35409_2.png) [@Zentrik](https://discourse.julialang.org/u/Zentrik)\
**Post date:** [May 28, 2025, 8:08pm UTC](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552/20 "2025-05-28T20:08:59Z")

</div>

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`.

[Next page](https://discourse.julialang.org/t/calling-python-function-jit-compiled-with-jax-from-julia-without-overhead/123552.md?page=2)