feat: export reactant compiled functions as tf SavedModel (#1426)

* fix: don't initialize jax arrays for tracing

* fix: disable tf gpus

* chore: setup code for serialization

* feat: mostly working

* feat: working saved model export

* chore: run fmt

* docs: serialization to savedmodel docs

* chore: fmt

* test: serialization

* Update Project.toml

Author: avik-pal

CondaPkg.toml

@@ -1,2 +1,4 @@
 [pip.deps]
-jax = ">=0.4"
+jax = ">= 0.6"
+tensorflow = ">= 2.17"
+numpy = ">= 2"

Project.toml

@@ -86,7 +86,7 @@ OneHotArrays = "0.2.10"
 OrderedCollections = "1"
 PrecompileTools = "1.2"
 Preferences = "1.4"
-PythonCall = "0.9"
+PythonCall = "0.9.25"
 Random = "1.10"
 Random123 = "1.7"
 ReactantCore = "0.1.15"

docs/src/.vitepress/config.mts

@@ -95,6 +95,7 @@ export default defineConfig({
         items: [
           { text: "Core Reactant API", link: "/api/api" },
           { text: "Sharding", link: "/api/sharding" },
+          { text: "Serialization", link: "/api/serialization" },
           { text: "Ops", link: "/api/ops" },
           { text: "Configuration", link: "/api/config" },
           {
@@ -169,6 +170,7 @@ export default defineConfig({
             link: "/api/api",
           },
           { text: "Sharding", link: "/api/sharding" },
+          { text: "Serialization", link: "/api/serialization" },
           { text: "Ops", link: "/api/ops" },
           { text: "Configuration", link: "/api/config" },
           {

docs/src/api/serialization.md

@@ -0,0 +1,29 @@
+```@meta
+CollapsedDocStrings = true
+```
+
+# Serialization
+
+```@docs
+Reactant.Serialization
+```
+
+## Exporting to TensorFlow SavedModel
+
+!!! note "Load PythonCall"
+
+    Serialization to TensorFlow SavedModel requires PythonCall to be loaded. Loading
+    PythonCall will automatically install tensorflow. If tensorflow installation fails,
+    we won't be able to export to SavedModel.
+
+A SavedModel contains a complete TensorFlow program, including trained parameters (i.e,
+tf.Variables) and computation. It does not require the original model building code to run,
+which makes it useful for sharing or deploying with [TFLite](https://tensorflow.org/lite),
+[TensorFlow.js](https://js.tensorflow.org/),
+[TensorFlow Serving](https://www.tensorflow.org/tfx/serving/tutorials/Serving_REST_simple),
+or [TensorFlow Hub](https://tensorflow.org/hub). Refer to the
+[official documentation](https://www.tensorflow.org/guide/saved_model) for more details.
+
+```@docs
+Reactant.Serialization.export_as_tf_saved_model
+```

ext/ReactantPythonCallExt.jl

@@ -0,0 +1,62 @@
+module ReactantPythonCallExt
+
+using PythonCall
+using Reactant: Reactant, TracedRArray
+
+const jaxptr = Ref{Py}()
+const jnpptr = Ref{Py}()
+
+const JAX_TRACING_SUPPORTED = Ref{Bool}(false)
+
+const tfptr = Ref{Py}()
+const tf2xlaptr = Ref{Py}()
+const npptr = Ref{Py}()
+
+const SAVED_MODEL_EXPORT_SUPPORTED = Ref{Bool}(false)
+
+const NUMPY_SIMPLE_TYPES = Dict(
+    Bool => :bool,
+    Int8 => :int8,
+    Int16 => :int16,
+    Int32 => :int32,
+    Int64 => :int64,
+    UInt8 => :uint8,
+    UInt16 => :uint16,
+    UInt32 => :uint32,
+    UInt64 => :uint64,
+    Float16 => :float16,
+    Float32 => :float32,
+    Float64 => :float64,
+    ComplexF16 => :complex16,
+    ComplexF32 => :complex32,
+    ComplexF64 => :complex64,
+)
+
+function __init__()
+    try
+        jaxptr[] = pyimport("jax")
+        jnpptr[] = pyimport("jax.numpy")
+        JAX_TRACING_SUPPORTED[] = true
+    catch err
+        @warn "Failed to import jax. Tracing jax functions invoked with pycall won't \
+               be supported." exception = (err, catch_backtrace())
+    end
+
+    try
+        tfptr[] = pyimport("tensorflow")
+        tfptr[].config.set_visible_devices(pylist(); device_type="GPU")
+        tf2xlaptr[] = pyimport("tensorflow.compiler.tf2xla.python.xla")
+        npptr[] = pyimport("numpy")
+        SAVED_MODEL_EXPORT_SUPPORTED[] = true
+    catch err
+        @warn "Failed to import tensorflow. Exporting Reactant compiled functions as \
+               tensorflow SavedModel will not be \
+               supported." exception = (err, catch_backtrace())
+    end
+    return nothing
+end
+
+include("pycall.jl")
+include("saved_model.jl")
+
+end

ext/ReactantPythonCallExt/ReactantPythonCallExt.jl

@@ -0,0 +1,18 @@
+function PythonCall.pycall(f::Py, arg0::TracedRArray, argNs::TracedRArray...; kwargs...)
+    JAX_TRACING_SUPPORTED[] || throw("jax could not be loaded.")
+
+    jax = jaxptr[]
+    jnp = jnpptr[]
+
+    inputs = map((arg0, argNs...)) do arg
+        jax.ShapeDtypeStruct(
+            size(arg),
+            jnp.dtype(string(NUMPY_SIMPLE_TYPES[Reactant.unwrapped_eltype(arg)])),
+        )
+    end
+
+    lowered = jax.jit(f).lower(inputs...)
+    res = Reactant.Ops.hlo_call(pyconvert(String, lowered.as_text()), arg0, argNs...)
+
+    return length(res) == 0 ? nothing : res[1]
+end

ext/ReactantPythonCallExt/pycall.jl

@@ -0,0 +1,92 @@
+# TODO: at some point, we should use the TF C++ API to export the SavedModel
+
+function Reactant.Serialization.serialization_supported(::Val{:SavedModel})
+    return SAVED_MODEL_EXPORT_SUPPORTED[]
+end
+
+function _extract_call_parameters(args::Tuple, input_locations, state_dict)
+    call_args = pylist([])
+    for loc in input_locations
+        if loc isa Reactant.Serialization.TFSavedModel.InputArgument
+            call_args.append(args[loc.position])
+        else
+            @assert haskey(state_dict, loc.name) "State dictionary does not contain key: \
+                                                  $(loc.name)"
+            call_args.append(state_dict[loc.name])
+        end
+    end
+    return call_args
+end
+
+function _wrap_as_tf_func(
+    spec::Reactant.Serialization.TFSavedModel.ReactantFunctionSpec, state_dict
+)
+    Touts = pylist([string(sig.dtype) for sig in spec.output_signature])
+    Souts = pylist([pylist(sig.shape) for sig in spec.output_signature])
+    return pyfunc(
+        function (args...)
+            return tf2xlaptr[].call_module(
+                _extract_call_parameters(args, spec.input_locations, state_dict);
+                version=5,
+                Tout=Touts,  # dtype information
+                Sout=Souts,  # Shape information
+                function_list=pylist([]),  # No functions to call
+                :module => spec.bytecode,
+            )
+        end,
+    )
+end
+
+function _make_input_signatures(
+    fn_spec::Reactant.Serialization.TFSavedModel.ReactantFunctionSpec
+)
+    input_pos_to_spec = Dict(
+        loc.position => spec for
+        (loc, spec) in zip(fn_spec.input_locations, fn_spec.input_signature) if
+        loc isa Reactant.Serialization.TFSavedModel.InputArgument
+    )
+
+    sigs = []
+    for i in 1:length(input_pos_to_spec)
+        spec = input_pos_to_spec[i]
+        dtype = getproperty(tfptr[], spec.dtype)
+        push!(
+            sigs,
+            tfptr[].TensorSpec(;
+                shape=pylist(spec.shape), dtype=dtype, name="args_$(i - 1)"
+            ),
+        )
+    end
+    return pylist(sigs)
+end
+
+function Reactant.Serialization.TFSavedModel.__to_tf_saved_model(
+    fn_spec::Reactant.Serialization.TFSavedModel.ReactantFunctionSpec, path::String
+)
+    tfm = tfptr[].Module()
+
+    state_dict = Dict(
+        k => tfptr[].Variable(
+            npptr[].asarray(permutedims(v, collect(ndims(v):-1:1)));
+            trainable=false,
+            name=k,
+        ) for (k, v) in fn_spec.state_dict
+    )
+
+    input_signatures = _make_input_signatures(fn_spec)
+
+    tfm.f = getproperty(tfptr[], :function)(
+        _wrap_as_tf_func(fn_spec, state_dict); input_signature=input_signatures
+    )
+    tfm._variables = pylist(collect(values(state_dict)))
+
+    signatures = pydict([
+        tfptr[].saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY =>
+            tfm.f.get_concrete_function(input_signatures...),
+    ])
+    save_options = tfptr[].saved_model.SaveOptions(; function_aliases=pydict(["" => tfm.f]))
+
+    tfptr[].saved_model.save(tfm, path; signatures=signatures, options=save_options)
+
+    return nothing
+end

ext/ReactantPythonCallExt/saved_model.jl

@@ -2449,11 +2449,16 @@ Compile the function `f` with arguments `args` and return the compiled function.
 
 $(SYNC_DOCS)
 $(COMMON_COMPILE_OPTIONS_DOCS)
+  - `serializable`: If `true`, the compiled function will be serializable. This is needed
+    for saving the compiled function to disk and loading it later. Defaults to `false`.
 
 See also [`@jit`](@ref), [`@code_hlo`](@ref), [`@code_mhlo`](@ref), [`@code_xla`](@ref).
 """
 macro compile(args...)
-    default_options = merge(get_common_compile_options(), Dict{Symbol,Any}(:sync => false))
+    default_options = merge(
+        get_common_compile_options(),
+        Dict{Symbol,Any}(:sync => false, :serializable => false),
+    )
     return esc(first(compile_call_expr(__module__, compile, default_options, args...)))
 end
 
@@ -3413,7 +3418,7 @@ function compile_xla(f, args; client=nothing, serializable::Bool=false, kwargs..
         # XLA.compile mutates the module, for serialization we need to keep a copy
         if serializable
             iobuffer = IOBuffer()
-            show(IOContext(iobuffer, :debug => debug), mod)
+            show(IOContext(iobuffer, :debug => true), mod)
             module_string = String(take!(iobuffer))
         else
             module_string = ""

src/Compiler.jl

@@ -190,6 +190,9 @@ include("Compiler.jl")
 
 include("Overlay.jl")
 
+# Serialization
+include("serialization/Serialization.jl")
+
 using .Compiler: @compile, @code_hlo, @code_mhlo, @jit, @code_xla, traced_getfield, compile
 export ConcreteRArray,
     ConcreteRNumber,