Support tf2jax.convert inside of tf.functions (e.g. when a JAX model is being traced for export as a TF saved model via orbax.export).#258

Closed

copybara-service[bot] wants to merge 1 commit intomainfrom

Commits on Oct 30, 2025

Support tf2jax.convert inside of tf.functions (e.g. when a JAX model is being traced for export as a TF saved model via orbax.export).

rryan
authored and
TF2JAXDev
committed