feat: generalize pycall to support structures (#1701)

* feat: generalize pycall to support structures

* chore: update test/integration/python.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: avik-pal

ext/ReactantPythonCallExt/ReactantPythonCallExt.jl

@@ -1,7 +1,7 @@
 module ReactantPythonCallExt
 
 using PythonCall: PythonCall, Py, pyconvert, pydict, pyfunc, pyimport, pylist
-using Reactant: Reactant, TracedRArray
+using Reactant: Reactant, TracedRArray, TracedRNumber, @reactant_overlay
 using Reactant.Ops: @opcall
 
 const jaxptr = Ref{Py}()
@@ -57,6 +57,7 @@ function __init__()
     return nothing
 end
 
+include("overlays.jl")
 include("pycall.jl")
 include("saved_model.jl")
 

ext/ReactantPythonCallExt/overlays.jl

@@ -0,0 +1,7 @@
+@reactant_overlay function PythonCall.pycall(f::Py, args...)
+    if Reactant.looped_any(Reactant.use_overlayed_version, args)
+        return pycall_with_jax_tracing(f, args...)
+    else
+        return Base.inferencebarrier(PythonCall.pycall)(f, args...)
+    end
+end

ext/ReactantPythonCallExt/pycall.jl

@@ -1,18 +1,37 @@
-function PythonCall.pycall(f::Py, arg0::TracedRArray, argNs::TracedRArray...; kwargs...)
+Reactant.jax_dtype_struct_type(::Type{T}) where {T} = Py
+
+function Reactant.convert_to_jax_dtype_struct(x::Union{TracedRArray,TracedRNumber})
     JAX_TRACING_SUPPORTED[] || throw("jax could not be loaded.")
+    return jaxptr[].ShapeDtypeStruct(
+        size(x), jnpptr[].dtype(string(NUMPY_SIMPLE_TYPES[Reactant.unwrapped_eltype(x)]))
+    )
+end
 
-    jax = jaxptr[]
-    jnp = jnpptr[]
+function pycall_with_jax_tracing(f::Py, args...)
+    JAX_TRACING_SUPPORTED[] || throw("jax could not be loaded.")
 
-    inputs = map((arg0, argNs...)) do arg
-        jax.ShapeDtypeStruct(
-            size(arg),
-            jnp.dtype(string(NUMPY_SIMPLE_TYPES[Reactant.unwrapped_eltype(arg)])),
-        )
+    seen_args = Reactant.OrderedIdDict()
+    jax_inputs = Vector{Any}(undef, length(args))
+    static_argnums = ()
+    prev_len = 0
+    for (i, arg) in enumerate(args)
+        jax_inputs[i] = Reactant.make_tracer(seen_args, arg, (), Reactant.TracedToJAX)
+        if length(seen_args) == prev_len
+            static_argnums = (static_argnums..., i - 1)
+        end
+        prev_len = length(seen_args)
     end
 
-    lowered = jax.jit(f).lower(inputs...)
-    res = @opcall hlo_call(pyconvert(String, lowered.as_text()), arg0, argNs...)
+    linear_args = Reactant.TracedType[]
+    for (k, v) in seen_args
+        k isa Reactant.TracedType || continue
+        push!(linear_args, k)
+    end
 
-    return length(res) == 0 ? nothing : res[1]
+    lowered = jaxptr[].jit(f; static_argnums).lower(jax_inputs...)
+    # To figure out the exact structure of the pyfunc, we need to execute it. Currently,
+    # we skip doing that and assume that we are returning nothing, array, or tuple of
+    # arrays.
+    res = @opcall hlo_call(pyconvert(String, lowered.as_text()), linear_args...)
+    return length(res) == 0 ? nothing : (length(res) == 1 ? res[1] : res)
 end

src/Tracing.jl

@@ -6,8 +6,12 @@
     TracedSetPath = 5
     TracedToTypes = 6
     NoStopTracedTrack = 7
+    TracedToJAX = 8
 end
 
+function convert_to_jax_dtype_struct end
+function jax_dtype_struct_type end
+
 struct VisitedObject
     id::Int
 end
@@ -386,6 +390,8 @@ Base.@nospecializeinfer function traced_type_inner(
             }
         end
         error("Unsupported runtime $runtime")
+    elseif mode == TracedToJAX
+        return jax_dtype_struct_type(T)
     elseif mode == TracedTrack || mode == NoStopTracedTrack || mode == TracedSetPath
         return T
     else
@@ -432,6 +438,8 @@ Base.@nospecializeinfer function traced_type_inner(
             }
         end
         error("Unsupported runtime $runtime")
+    elseif mode == TracedToJAX
+        return jax_dtype_struct_type(T)
     elseif mode == TracedTrack || mode == NoStopTracedTrack || mode == TracedSetPath
         return T
     else
@@ -1312,6 +1320,17 @@ Base.@nospecializeinfer function make_tracer(
         error("Unsupported runtime $runtime")
     end
 
+    if mode == TracedToJAX
+        haskey(seen, prev) && return seen[prev]
+        if !Sharding.is_sharded(sharding)
+            res = convert_to_jax_dtype_struct(prev)
+        else
+            error("TODO: implement sharding")
+        end
+        seen[prev] = res
+        return res
+    end
+
     throw("Cannot Unknown trace mode $mode")
 end
 
@@ -1390,6 +1409,17 @@ Base.@nospecializeinfer function make_tracer(
         error("Unsupported runtime $runtime")
     end
 
+    if mode == TracedToJAX
+        haskey(seen, prev) && return seen[prev]
+        if !Sharding.is_sharded(sharding)
+            res = convert_to_jax_dtype_struct(prev)
+        else
+            error("TODO: implement sharding")
+        end
+        seen[prev] = res
+        return res
+    end
+
     throw("Cannot Unknown trace mode $mode")
 end
 

test/integration/python.jl

@@ -11,10 +11,33 @@ fn(x, y) = sin.(x) .+ cos.(y.x[1:2, :])
 
     @testset "PythonCall" begin
         jax = pyimport("jax")
+        jnp = pyimport("jax.numpy")
 
         result = @jit jax.numpy.sum(Reactant.to_rarray(Float32[1, 2, 3]))
         @test result isa ConcreteRNumber{Float32}
         @test result ≈ 6
+
+        pyfn = pyfunc(
+            function (tup, y, num, partial_eval)
+                return jnp.sin(tup[0]) + jnp.cos(tup[1]) + jnp.sum(y) - num * partial_eval
+            end;
+            name="sum_and_sin_cos_jax",
+        )
+
+        tup = (rand(Float32, 3, 4), rand(Float32, 3, 4))
+        y = rand(Float32, 2, 2)
+        num = 4.0
+        partial_eval = 0.5
+
+        gt = sin.(tup[1]) .+ cos.(tup[2]) .+ sum(y) .- num .* partial_eval
+
+        tup_ra = Reactant.to_rarray(tup)
+        y_ra = Reactant.to_rarray(y)
+        num_ra = ConcreteRNumber{Float32}(num)
+
+        result = @jit pyfn(tup_ra, y_ra, num_ra, partial_eval)
+        @test result ≈ gt atol = 1e-5
+        @test result isa ConcreteRArray{Float32,2}
     end
 
     @testset "SavedModel Export" begin