feat: add size checks

Author: avik-pal

docs/src/api/serialization.md

@@ -30,15 +30,14 @@ Reactant.Serialization.export_as_tf_saved_model
 
 ## Exporting to JAX via EnzymeAD
 
-!!! note "No Dependencies Required"
+!!! note "Load NPZ"
 
-    Unlike TensorFlow SavedModel export, exporting to JAX via EnzymeAD does not require any
-    Python dependencies at build time. It generates standalone files that can be used with
-    EnzymeAD/JAX in Python.
+    This export functionality requires the `NPZ` package to be loaded.
 
 This export functionality generates:
+
 1. A `.mlir` file containing the StableHLO representation of your Julia function
-2. Example input `.npy` files with properly transposed arrays (column-major → row-major)
+2. Input `.npz` files containing the input arrays for the function
 3. A Python script that wraps the function for use with `enzyme_ad.jax.hlo_call`
 
 The generated Python script can be immediately used with JAX and EnzymeAD without any

src/serialization/EnzymeJAX.jl

@@ -33,6 +33,8 @@ This function:
   - `output_dir::Union{String,Nothing}`: Directory where output files will be saved. If
     `nothing`, uses a temporary directory and prints the path.
   - `function_name::String`: Base name for generated files
+  - `preserve_sharding::Bool`: Whether to preserve sharding information in the exported
+    function. Defaults to `true`.
 
 ## Returns
 
@@ -55,8 +57,11 @@ function my_function(x::AbstractArray, y::NamedTuple, z::Number)
 end
 
 # Create some example inputs
-x = Reactant.to_rarray(Float32[1, 2, 3])
-y = (; x=Reactant.to_rarray(Float32[4, 5, 6]), y=Reactant.to_rarray(Float32[7, 8, 9]))
+x = Reactant.to_rarray(reshape(collect(Float32, 1:6), 2, 3))
+y = (;
+    x=Reactant.to_rarray(reshape(collect(Float32, 7:12), 2, 3)),
+    y=Reactant.to_rarray(reshape(collect(Float32, 13:18), 2, 3))
+)
 z = Reactant.to_rarray(10.0f0; track_numbers=true)
 
 # Export to EnzymeJAX
@@ -73,7 +78,11 @@ result = jax.jit(run_my_function)(*inputs)
 ```
 """
 function export_to_enzymejax(
-    f, args...; output_dir::Union{String,Nothing}=nothing, function_name::String=string(f)
+    f,
+    args...;
+    output_dir::Union{String,Nothing}=nothing,
+    function_name::String=string(f),
+    preserve_sharding::Bool=true,
 )
     if output_dir === nothing
         output_dir = mktempdir(; cleanup=false)
@@ -173,6 +182,15 @@ function _generate_python_script(
         "\n",
     )
 
+    arg_size_checks = [
+        "assert $(arg_names[i]).shape == $(reverse(info.shape)), f\"Expected shape of $(arg_names[i]) to be $(reverse(info.shape)). Got {$(arg_names[i]).shape} (path: $(info.path))\""
+        for (i, info) in enumerate(input_info)
+    ]
+    arg_dtype_checks = [
+        "assert $(arg_names[i]).dtype == np.dtype('$(Serialization.NUMPY_SIMPLE_TYPES[info.dtype])'), f\"Expected dtype of $(arg_names[i]) to be $(Serialization.NUMPY_SIMPLE_TYPES[info.dtype]). Got {$(arg_names[i]).dtype} (path: $(info.path))\""
+        for (i, info) in enumerate(input_info)
+    ]
+
     load_inputs = ["npz_data['$(info.key)']" for info in input_info]
 
     # Build the complete Python script
@@ -217,6 +235,8 @@ function _generate_python_script(
             arrays from Julia, make sure to transpose them first using:
             \`permutedims(arr, reverse(1:ndims(arr)))\`
         \"\"\"
+        $(join(arg_dtype_checks, "\n    "))
+        $(join(arg_size_checks, "\n    "))
         return hlo_call(
             $(arg_list),
             source=_hlo_code,

src/serialization/Serialization.jl

@@ -127,20 +127,6 @@ function export_as_tf_saved_model(
     )
 end
 
-"""
-    export_to_enzymejax(
-        f,
-        args...;
-        output_dir::String=".",
-        function_name::String="exported_function",
-    )
-
-Export a Julia function to EnzymeJAX format for use in Python/JAX.
-
-See [`EnzymeJAX.export_to_enzymejax`](@ref) for details.
-"""
-function export_to_enzymejax(f, args...; kwargs...)
-    return EnzymeJAX.export_to_enzymejax(f, args...; kwargs...)
-end
+const export_to_enzymejax = EnzymeJAX.export_to_enzymejax
 
 end