Pythoncall (#407)

* WIP: pythoncall

* fix

* fix

* Update Project.toml

* Apply suggestions from code review

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

* fix

---------

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

Author: wsmoses

CondaPkg.toml

@@ -0,0 +1,2 @@
+[deps]
+jax = ""

Project.toml

@@ -24,18 +24,17 @@ AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 Random123 = "74087812-796a-5b5d-8853-05524746bad3"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 YaoBlocks = "418bc28f-b43b-5e0b-a6e7-61bbc1a2c1df"
 
-[sources]
-ReactantCore = {path = "lib/ReactantCore"}
-
 [extensions]
 ReactantAbstractFFTsExt = "AbstractFFTs"
 ReactantArrayInterfaceExt = "ArrayInterface"
 ReactantCUDAExt = "CUDA"
 ReactantNNlibExt = "NNlib"
+ReactantPythonCallExt = "PythonCall"
 ReactantRandom123Ext = "Random123"
 ReactantStatisticsExt = "Statistics"
 ReactantYaoBlocksExt = "YaoBlocks"
@@ -68,3 +67,7 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
+
+[sources.ReactantCore]
+path = "lib/ReactantCore"

ext/ReactantPythonCallExt.jl

@@ -0,0 +1,41 @@
+module ReactantPythonCallExt
+
+using PythonCall
+using Reactant: Reactant, TracedRArray, AnyTracedRArray, MLIR, TracedRNumber
+using ReactantCore: @trace
+
+using PythonCall
+
+const jaxptr = Ref{Py}()
+
+function PythonCall.pycall(
+    f::Py, arg0::Reactant.TracedRArray, argNs::Reactant.TracedRArray...; kwargs...
+)
+    jax = jaxptr[]
+    numpy = jax.numpy
+    inputs = map((arg0, argNs...)) do arg
+        JT = eltype(arg)
+        PT = nothing
+        for (CPT, CJT) in PythonCall.Convert.NUMPY_SIMPLE_TYPES
+            if JT == CJT
+                PT = CPT
+                break
+            end
+        end
+        numpy.zeros(size(arg); dtype=getproperty(numpy, Symbol(PT)))
+    end
+    lowered = jax.jit(f).lower(inputs...)
+    txt = pyconvert(String, lowered.as_text())
+    res = Reactant.Ops.hlo_call(txt, arg0, argNs...)
+    if length(res) == 0
+        return nothing
+    else
+        return res[1]
+    end
+end
+
+function __init__()
+    return jaxptr[] = pyimport("jax")
+end
+
+end # module ReactantPythonCallExt

test/Project.toml

@@ -16,6 +16,7 @@ MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Random123 = "74087812-796a-5b5d-8853-05524746bad3"
 Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"

test/python.jl

@@ -0,0 +1,13 @@
+using Reactant
+using Reactant: Ops
+
+using Test
+using PythonCall
+
+@testset "PythonCall" begin
+    jax = pyimport("jax")
+
+    result = @jit jax.numpy.sum(Reactant.to_rarray(Float32[1, 2, 3]))
+    @test typeof(result) == ConcreteRNumber{Float32}
+    @test result ≈ 6
+end

test/runtests.jl

@@ -56,6 +56,7 @@ const REACTANT_TEST_GROUP = lowercase(get(ENV, "REACTANT_TEST_GROUP", "all"))
         @safetestset "Shortcuts to MLIR ops" include("ops.jl")
         @safetestset "Wrapped Arrays" include("wrapped_arrays.jl")
         @safetestset "Control Flow" include("control_flow.jl")
+        @safetestset "Python" include("python.jl")
     end
 
     if REACTANT_TEST_GROUP == "all" || REACTANT_TEST_GROUP == "integration"