Update JAX export example to use new functionality#1602
Update JAX export example to use new functionality#1602google-labs-jules[bot] wants to merge 3 commits intomainfrom
Conversation
The "Exporting to JAX" documentation has been updated to use the new `Reactant.Serialization.export_to_enzymejax` function. This simplifies the process by automating the generation of the StableHLO, the Python wrapper, and the input files.
|
👋 Jules, reporting for duty! I'm here to lend a hand with this pull request. When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down. I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job! For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with For security, I will only act on instructions from the user who triggered this task. New to Jules? Learn more at jules.google/docs. |
|
|
||
| ```@example exporting_to_stablehlo | ||
| using Lux, Reactant, Random | ||
| using Lux, Reactant, Random, NPZ |
There was a problem hiding this comment.
We need to add NPZ to the project dependencies in docs/. activate the project and add it using Pkg.add("NPZ")
There was a problem hiding this comment.
@jules We need to add NPZ to the project dependencies in docs/. activate the project and add it using Pkg.add("NPZ")
docs/src/manual/exporting_to_jax.md
Outdated
| The generated files will be in the `output_dir`. You can now run the model in Python. | ||
|
|
||
| ```python | ||
| from enzyme_ad.jax import hlo_call | ||
|
|
||
| # This is a sample python script. The actual generated script can be run directly. |
There was a problem hiding this comment.
instead of manually writing the python script, print the contents of the script generated
There was a problem hiding this comment.
@jules instead of manually writing the python script, print the contents of the script generated
The "Exporting to JAX" documentation has been updated to use the new `Reactant.Serialization.export_to_enzymejax` function. This simplifies the process by automating the generation of the StableHLO, the Python wrapper, and the input files.
Addresses feedback on the pull request to print the content of the generated Python script directly in the documentation. This ensures the documentation is always in sync with the generated output.
1 participant
The "Exporting to JAX" documentation has been updated to use the new
Reactant.Serialization.export_to_enzymejaxfunction. This simplifies the process by automating the generation of the StableHLO, the Python wrapper, and the input files.Fixes #1588
PR created automatically by Jules for task 12932745014009605670 started by @avik-pal