Skip to content

Commit 3bd8d73

Browse files
Copilotavik-pal
andcommitted
Add documentation for export_to_enzymeax function
Co-authored-by: avik-pal <[email protected]>
1 parent 7e24c31 commit 3bd8d73

File tree

1 file changed

+20
-0
lines changed

1 file changed

+20
-0
lines changed

docs/src/api/serialization.md

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,3 +27,23 @@ or [TensorFlow Hub](https://tensorflow.org/hub). Refer to the
2727
```@docs
2828
Reactant.Serialization.export_as_tf_saved_model
2929
```
30+
31+
## Exporting to JAX via EnzymeAD
32+
33+
!!! note "No Dependencies Required"
34+
35+
Unlike TensorFlow SavedModel export, exporting to JAX via EnzymeAD does not require any
36+
Python dependencies at build time. It generates standalone files that can be used with
37+
EnzymeAD/JAX in Python.
38+
39+
This export functionality generates:
40+
1. A `.mlir` file containing the StableHLO representation of your Julia function
41+
2. Example input `.npy` files with properly transposed arrays (column-major → row-major)
42+
3. A Python script that wraps the function for use with `enzyme_ad.jax.hlo_call`
43+
44+
The generated Python script can be immediately used with JAX and EnzymeAD without any
45+
additional Julia dependencies.
46+
47+
```@docs
48+
Reactant.Serialization.export_to_enzymeax
49+
```

0 commit comments

Comments
 (0)