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

A while ago we moved the native exporting code out of jax2tf, to
jax.experimental.export, but we left behind shape_poly.py shim for
backwards compatibility. Now we remove this shim.

Replace previous

  from jax.experimental.jax2tf import shape_poly

with

  from jax.experimental.export import shape_poly

PiperOrigin-RevId: 589850402

jax/experimental/jax2tf/BUILD

@@ -41,7 +41,6 @@ py_library(
         "call_tf.py",
         "impl_no_xla.py",
         "jax2tf.py",
-        "shape_poly.py",  # TODO(necula): remove stub
     ],
     srcs_version = "PY3",
     # TODO: b/255503696: enable pytype

jax/experimental/jax2tf/__init__.py

@@ -21,5 +21,3 @@
   PolyShape as PolyShape
 )
 from jax.experimental.jax2tf.call_tf import call_tf as call_tf
-# Needed by maths.qec.
-from jax.experimental.jax2tf import shape_poly

jax/experimental/jax2tf/tests/models_test_main.py

@@ -54,7 +54,7 @@
 
 import tensorflow as tf
 
-from jax.experimental.jax2tf.shape_poly import InconclusiveDimensionOperation
+from jax.experimental.export.shape_poly import InconclusiveDimensionOperation
 from jax.experimental.jax2tf.tests.model_harness import ALL_HARNESSES
 from jax.experimental.jax2tf.tests.converters import ALL_CONVERTERS