keras/src/export/export_lib.py/ExportArchive/_convert_jax2tf_function
class ExportArchive:    def _convert_jax2tf_function(self, fn, input_signature, jax2tf_kwargs=None):
        from jax.experimental import jax2tf

        if jax2tf_kwargs is None:
            jax2tf_kwargs = {}

        if "native_serialization" not in jax2tf_kwargs:
            jax2tf_kwargs["native_serialization"] = (
                self._check_device_compatible()
            )

        variables_shapes = self._to_polymorphic_shape(
            self._backend_variables, allow_none=False
        )
        if "polymorphic_shapes" in jax2tf_kwargs:
            input_shapes = jax2tf_kwargs["polymorphic_shapes"]
        else:
            input_shapes = self._to_polymorphic_shape(input_signature)
        jax2tf_kwargs["polymorphic_shapes"] = [variables_shapes] + input_shapes

        return jax2tf.convert(fn, **jax2tf_kwargs)