Fix code review issues: binary mode, ComplexF16 support, Python indentation

Co-authored-by: avik-pal <[email protected]>

Author: Copilot

src/serialization/EnzymeJAX.jl

@@ -15,7 +15,6 @@ const NUMPY_SIMPLE_TYPES = Dict(
     Float16 => "np.float16",
     Float32 => "np.float32",
     Float64 => "np.float64",
-    ComplexF16 => "np.complex64",  # Note: NumPy doesn't have float16 complex
     ComplexF32 => "np.complex64",
     ComplexF64 => "np.complex128",
 )
@@ -177,10 +176,11 @@ function _save_transposed_array(path::String, arr::AbstractArray)
         shape_str = join(size(transposed), ", ")
         header = "{'descr': '$(dtype_str)', 'fortran_order': False, 'shape': ($(shape_str),)}"
 
-        # Pad header to be aligned on 64 bytes
+        # Pad header to be aligned on 64 bytes (16-byte alignment for v1.0)
+        # Total size needs to be divisible by 16
         header_len = length(header) + 1  # +1 for newline
         total_len = 10 + header_len  # 10 = magic(6) + version(2) + header_len(2)
-        padding = (64 - (total_len % 64)) % 64
+        padding = (16 - (total_len % 16)) % 16
         header = header * " "^padding * "\n"
         header_len = length(header)
 
@@ -191,6 +191,7 @@ function _save_transposed_array(path::String, arr::AbstractArray)
         # Write data
         write(io, vec(transposed))
     end
+    return nothing
 end
 
 """
@@ -267,99 +268,86 @@ function _generate_python_script(
     mlir_rel = relpath(mlir_path, output_dir)
     input_rels = [relpath(p, output_dir) for p in input_paths]
 
-    # Start building the Python script
-    script = """
-    \"\"\"
-    Auto-generated Python script for calling exported Julia/Reactant function via EnzymeJAX.
+    # Build the Python script without leading indentation
+    lines = String[]
 
-    This script was generated by Reactant.Serialization.export_to_enzymeax().
-    \"\"\"
+    # Header
+    push!(lines, "\"\"\"")
+    push!(lines, "Auto-generated Python script for calling exported Julia/Reactant function via EnzymeJAX.")
+    push!(lines, "")
+    push!(lines, "This script was generated by Reactant.Serialization.export_to_enzymeax().")
+    push!(lines, "\"\"\"")
+    push!(lines, "")
+    push!(lines, "from enzyme_ad.jax import hlo_call")
+    push!(lines, "import jax")
+    push!(lines, "import jax.numpy as jnp")
+    push!(lines, "import numpy as np")
+    push!(lines, "import os")
+    push!(lines, "")
+    push!(lines, "# Get the directory of this script")
+    push!(lines, "_script_dir = os.path.dirname(os.path.abspath(__file__))")
+    push!(lines, "")
+    push!(lines, "# Load the MLIR/StableHLO code")
+    push!(lines, "with open(os.path.join(_script_dir, \"$(mlir_rel)\"), \"r\") as f:")
+    push!(lines, "    _hlo_code = f.read()")
+    push!(lines, "")
 
-    from enzyme_ad.jax import hlo_call
-    import jax
-    import jax.numpy as jnp
-    import numpy as np
-    import os
-    
-    # Get the directory of this script
-    _script_dir = os.path.dirname(os.path.abspath(__file__))
-    
-    # Load the MLIR/StableHLO code
-    with open(os.path.join(_script_dir, "$(mlir_rel)"), "r") as f:
-        _hlo_code = f.read()
-    
-    """
-    
-    # Add function to load inputs
-    script *= """
-    def load_inputs():
-        \"\"\"Load the example inputs that were exported from Julia.\"\"\"
-        inputs = []
-    """
-    
-    for (i, input_rel) in enumerate(input_rels)
-        script *= """
-            inputs.append(np.load(os.path.join(_script_dir, "$(input_rel)")))
-        """
+    # Function to load inputs
+    push!(lines, "def load_inputs():")
+    push!(lines, "    \"\"\"Load the example inputs that were exported from Julia.\"\"\"")
+    push!(lines, "    inputs = []")
+    for input_rel in input_rels
+        push!(lines, "    inputs.append(np.load(os.path.join(_script_dir, \"$(input_rel)\")))")
     end
+    push!(lines, "    return tuple(inputs)")
+    push!(lines, "")
 
-    script *= """
-        return tuple(inputs)
-    
-    """
-    
-    # Add the main function that calls the HLO code
+    # Main function
     arg_names = ["arg$i" for i in 1:length(input_paths)]
     arg_list = join(arg_names, ", ")
 
-    script *= """
-    def run_$(function_name)($(arg_list)):
-        \"\"\"
-        Call the exported Julia function via EnzymeJAX.
-        
-        Args:
-    """
+    push!(lines, "def run_$(function_name)($(arg_list)):")
+    push!(lines, "    \"\"\"")
+    push!(lines, "    Call the exported Julia function via EnzymeJAX.")
+    push!(lines, "    ")
+    push!(lines, "    Args:")
 
     for (i, info) in enumerate(input_info)
         # Note: shapes are already transposed for Python
         python_shape = reverse(info.shape)
-        script *= """
-                $(arg_names[i]): Array of shape $(python_shape) and dtype $(NUMPY_SIMPLE_TYPES[info.dtype])
-        """
+        push!(lines, "        $(arg_names[i]): Array of shape $(python_shape) and dtype $(NUMPY_SIMPLE_TYPES[info.dtype])")
     end
 
-    script *= """
-        
-        Returns:
-            The result of calling the exported function.
-        
-        Note:
-            All inputs must be in row-major (Python/NumPy) order. If you're passing
-            arrays from Julia, make sure to transpose them first using:
-            `permutedims(arr, reverse(1:ndims(arr)))`
-        \"\"\"
-        return hlo_call(
-            $(arg_list),
-            source=_hlo_code,
-        )
-    
-    """
+    push!(lines, "    ")
+    push!(lines, "    Returns:")
+    push!(lines, "        The result of calling the exported function.")
+    push!(lines, "    ")
+    push!(lines, "    Note:")
+    push!(lines, "        All inputs must be in row-major (Python/NumPy) order. If you're passing")
+    push!(lines, "        arrays from Julia, make sure to transpose them first using:")
+    push!(lines, "        `permutedims(arr, reverse(1:ndims(arr)))`")
+    push!(lines, "    \"\"\"")
+    push!(lines, "    return hlo_call(")
+    push!(lines, "        $(arg_list),")
+    push!(lines, "        source=_hlo_code,")
+    push!(lines, "    )")
+    push!(lines, "")
 
-    # Add a main block for testing
-    script *= """
-    if __name__ == "__main__":
-        # Load the example inputs
-        inputs = load_inputs()
-        
-        # Run the function (with JIT compilation)
-        print("Running $(function_name) with JIT compilation...")
-        result = jax.jit(run_$(function_name))(*inputs)
-        print("Result:", result)
-        print("Result shape:", result.shape if hasattr(result, 'shape') else 'scalar')
-        print("Result dtype:", result.dtype if hasattr(result, 'dtype') else type(result))
-    """
+    # Main block
+    push!(lines, "if __name__ == \"__main__\":")
+    push!(lines, "    # Load the example inputs")
+    push!(lines, "    inputs = load_inputs()")
+    push!(lines, "    ")
+    push!(lines, "    # Run the function (with JIT compilation)")
+    push!(lines, "    print(\"Running $(function_name) with JIT compilation...\")")
+    push!(lines, "    result = jax.jit(run_$(function_name))(*inputs)")
+    push!(lines, "    print(\"Result:\", result)")
+    push!(lines, "    print(\"Result shape:\", result.shape if hasattr(result, 'shape') else 'scalar')")
+    push!(lines, "    print(\"Result dtype:\", result.dtype if hasattr(result, 'dtype') else type(result))")
 
-    write(python_path, script)
+    # Write the script
+    write(python_path, join(lines, "\n") * "\n")
+    return nothing
 end
 
 end  # module

test/runtests.jl

@@ -38,6 +38,7 @@ const REACTANT_TEST_GROUP = lowercase(get(ENV, "REACTANT_TEST_GROUP", "all"))
         @safetestset "Cluster Detection" include("cluster_detector.jl")
         @safetestset "Config" include("config.jl")
         @safetestset "Batching" include("batching.jl")
+        @safetestset "Export to EnzymeJAX" include("export_enzymeax.jl")
         @safetestset "QA" include("qa.jl")
     end