Internal changes (moving code around)

export/orbax/export/data_processors/data_processor_base.py

@@ -19,6 +19,7 @@
 import abc
 from collections.abc import Callable, Set
 from typing import Any
+import jaxtyping
 
 
 class DataProcessor(abc.ABC):

export/orbax/export/data_processors/jax_data_processor.py

@@ -24,16 +24,14 @@
 from orbax.export import obm_configs
 from orbax.export.data_processors import data_processor_base
 
-from .third_party.neptune.protos import manifest_pb2
-
 
 def _jax_spec_from(spec: Any) -> jax.ShapeDtypeStruct:
   """Converts a ShloTensorSpec to a jax.ShapeDtypeStruct."""
-  if isinstance(spec, shlo_function.ShloTensorSpec):
-    if spec.dtype == shlo_function.ShloDType.bf16:
+  if isinstance(spec, shlo_type.ShloTensorSpec):
+    if spec.dtype == shlo_type.ShloDType.bf16:
       return jax.ShapeDtypeStruct(spec.shape, jax.numpy.bfloat16)
     return jax.ShapeDtypeStruct(
-        spec.shape, shlo_function.shlo_dtype_to_np_dtype(spec.dtype)
+        spec.shape, shlo_type.shlo_dtype_to_np_dtype(spec.dtype)
     )
   if hasattr(spec, 'shape') and hasattr(spec, 'dtype'):
     return jax.ShapeDtypeStruct(

export/orbax/export/data_processors/jax_data_processor_test.py

@@ -19,12 +19,12 @@
 
 import jax
 import jax.numpy as jnp
-from orbax.experimental.model.core.python import value
 from orbax.export import obm_configs
 from orbax.export.data_processors import jax_data_processor
 
 from absl.testing import absltest
 from .testing.pybase import parameterized
+from .third_party.neptune.neptune_model._src.core import value
 from .third_party.neptune.protos import manifest_pb2
 
 

export/orbax/export/data_processors/tf_data_processor_test.py

@@ -13,10 +13,12 @@
 # limitations under the License.
 
 import os
-import orbax.experimental.model.core as obm
+
 from orbax.export.data_processors import tf_data_processor
 import tensorflow as tf
+
 from absl.testing import absltest
+import .third_party.neptune.neptune_model._src.core as obm
 from .util.task.python import error as google_error
 
 

export/orbax/export/export_testing_utils.py

@@ -27,6 +27,7 @@
 from orbax.export import jax_module
 from orbax.export import obm_configs
 from orbax.export import serving_config as osc
+
 import tensorflow as tf
 
 

export/orbax/export/obm_export_test.py

@@ -24,7 +24,7 @@
 from google.protobuf import text_format
 import jax
 import jax.numpy as jnp
-from jaxtyping import PyTree
+import jaxtyping
 from orbax.export import constants
 from orbax.export import export_testing_utils
 from orbax.export import jax_module

export/orbax/export/utils_test.py

@@ -21,7 +21,6 @@
 from orbax.export import utils
 import tensorflow as tf
 
-
 TensorSpecWithDefault = utils.TensorSpecWithDefault