mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #25872 from gnecula:jax2tf_doc

Browse Source 
PiperOrigin-RevId: 715411235
This commit is contained in:
jax authors 2025-01-14 09:24:20 -08:00
commit f270739f9f
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
jax/experimental/jax2tf/README.md
View File

@ -1864,12 +1864,12 @@ purposes of `call_tf`.)
Inside Google, you can turn on logging by using the `--vmodule` argument to
specify the logging levels for different modules,
e.g., `--vmodule=jax_export=3`. You can set `TF_DUMP_GRAPH_PREFIX` to
e.g., `--vmodule=_export=3`. You can set `TF_DUMP_GRAPH_PREFIX` to
a directory where modules should be dumped, or to `"-"` to dump the
modules to the log.
The following modules are useful for debugging JAX native serialization:
* `jax_export=3` - will log the StableHLO module on serialization.
* `_export=3` - will log the StableHLO module on serialization.
* `jax2tf=3` - will log the parameters to `XlaCallModule` op on serialization.
* `xla_call_module_loader=3` - will log the StableHLO module upon loading,
after shape refinements, and on verification error. You can use level `4` to