[jax2tf] Remove backwards compatibility shims (jax2tf.jax_export)

A while ago we moved the native exporting code out of jax2tf, to jax.experimental.export, but we left behind jax_export.py shim for backwards compatibility. Now we remove this shim. Replace previous from jax.experimental.jax2tf import jax_export with from jax.experimental.export import export PiperOrigin-RevId: 589409387
google · Dec 9, 2023 · e50ef1b · e50ef1b
1 parent 709564a
commit e50ef1b
Show file tree

Hide file tree

Showing 3 changed files with 0 additions and 25 deletions.
diff --git a/jax/experimental/jax2tf/BUILD b/jax/experimental/jax2tf/BUILD
@@ -41,7 +41,6 @@ py_library(
         "call_tf.py",
         "impl_no_xla.py",
         "jax2tf.py",
-        "jax_export.py",  # TODO(necula): remove stub
         "shape_poly.py",  # TODO(necula): remove stub
     ],
     srcs_version = "PY3",

diff --git a/jax/experimental/jax2tf/__init__.py b/jax/experimental/jax2tf/__init__.py
@@ -21,7 +21,5 @@
   PolyShape as PolyShape
 )
 from jax.experimental.jax2tf.call_tf import call_tf as call_tf
-# TODO(necula): remove stub. Needed by SAX
-from jax.experimental.jax2tf import jax_export
 # Needed by maths.qec.
 from jax.experimental.jax2tf import shape_poly
diff --git a/jax/experimental/jax2tf/jax_export.py b/jax/experimental/jax2tf/jax_export.py