jax.experimental.jax2tf with non-native serialization has been removed

[gnecula]

Hello JAX users,

If your code does not use jax.experimental.jax2tf, you can stop reading now.

The jax2tf API used to support two modes: the new one, called native serialization (native_serialization=True, the default since August 2023), uses the JAX native lowering rules to produce a StableHLO module that is then wrapped as a single TensorFlow op.
The older mode called graph or non-native serialization used to convert every JAX primitive to one or more TensorFlow ops.
See the documentation for a comparison and for compatibility recipes. In particular, if all you want is to execute the result of jax2tf then native serialization will give you a more accurate representation of the JAX program. If your code was processing the result of the conversion looking for particular TensorFlow ops, this is not possible anymore with native serialization.

In July 2024 (JAX version 0.4.31) we have deprecated the graph serialization and starting with the upcoming JAX version 0.4.36 only the native serialization mode will be supported.