[jax2tf] Remove most of the code for jax2tf with non-native serialization

We keep only the `enable_xla` and `native_serialization` parameters to `jax2tf.convert`. They are now deprecated and will be removed in a future version of JAX.

PiperOrigin-RevId: 800466991

Author: gnecula

CHANGELOG.md

@@ -29,6 +29,11 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
     `False`, but note this is a temporary config that will be removed in a
     future release.
 
+* Deprecations:
+  * The parameters `enable_xla` and `native_serialization` for `jax2tf.convert`
+    are deprecated and will be removed in a future versionof JAX. These were
+    used for jax2tf with non-native serialization, which has been now removed.
+
 ## JAX 0.7.1 (August 20, 2025)
 
 * New features

jax/experimental/jax2tf/BUILD

@@ -38,7 +38,6 @@ py_library(
     name = "jax2tf_internal",
     srcs = [
         "call_tf.py",
-        "impl_no_xla.py",
         "jax2tf.py",
     ],
     # TODO: b/255503696: enable pytype

jax/experimental/jax2tf/call_tf.py

@@ -303,10 +303,7 @@ def check_tf_result(idx: int, r_tf: TfVal, r_aval: core.ShapedArray | None) -> T
   # that tf.ensure_shape did this, but it can only take shapes that contain None
   # not computed shapes. However, in eager mode we should be able to resolve
   # the declared shapes to constants and we get better checking.
-  if tf.executing_eagerly():
-    r_aval_shape_tf = jax2tf_internal._eval_shape(r_aval.shape)
-  else:
-    r_aval_shape_tf = jax2tf_internal._aval_to_tf_shape(r_aval)
+  r_aval_shape_tf = jax2tf_internal._aval_to_tf_shape(r_aval)
   # We do as much checking as we can here, instead of relying on tf.ensure_shape
   # because the latter gives different errors in eager vs. compiled mode.
   # TODO(b/279454591): This strange error is from TF. Eager function suppose
@@ -646,16 +643,6 @@ def _register_call_lowering(platform):
 for platform in ("cpu", "cuda", "tpu"):
   _register_call_lowering(platform)
 
-# Support the call_tf under jax2tf.convert in eager mode
-def _jax2tf_call_tf(*args: TfVal,
-                    callable_flat_tf: Callable,
-                    **_) -> TfVal:
-  with jax2tf_internal.inside_call_tf():
-    res_tf_flat = callable_flat_tf(*args)
-  return res_tf_flat
-
-jax2tf_internal.tf_impl[call_tf_p] = _jax2tf_call_tf
-
 
 def emit_tf_embedded_graph_custom_call(
     ctx: mlir.LoweringRuleContext,

jax/experimental/jax2tf/examples/saved_model_lib.py

@@ -41,7 +41,6 @@ def convert_and_save_model(
     input_signatures: Sequence[tf.TensorSpec],
     polymorphic_shapes: str | None = None,
     with_gradient: bool = False,
-    enable_xla: bool = True,
     compile_model: bool = True,
     saved_model_options: tf.saved_model.SaveOptions | None = None):
   """Convert a JAX function and saves a SavedModel.
@@ -80,8 +79,6 @@ def convert_and_save_model(
       corresponding input shapes.
     with_gradient: the value to use for the `with_gradient` parameter for
       `jax2tf.convert`.
-    enable_xla: the value to use for the `enable_xla` parameter for
-      `jax2tf.convert`.
     compile_model: use TensorFlow jit_compiler on the SavedModel. This
       is needed if the SavedModel will be used for TensorFlow serving.
     polymorphic_shapes: if given then it will be used as the
@@ -99,8 +96,7 @@ def convert_and_save_model(
   tf_fn = jax2tf.convert(
     jax_fn,
     with_gradient=with_gradient,
-    polymorphic_shapes=[None, polymorphic_shapes],
-    enable_xla=enable_xla)
+    polymorphic_shapes=[None, polymorphic_shapes])
 
   # Create tf.Variables for the parameters. If you want more useful variable
   # names, you can use `tree.map_structure_with_path` from the `dm-tree` package