From 3cc828d4410c97429a0d46aa00adca7ac8b518f4 Mon Sep 17 00:00:00 2001 From: Orbax Authors Date: Mon, 27 Jul 2026 12:23:21 -0700 Subject: [PATCH] Fix Orbax export compatibility with JAX ShapeDtypeStruct validation. PiperOrigin-RevId: 954765777 --- .../export/data_processors/jax_data_processor.py | 11 ++++++----- export/orbax/export/typing.py | 6 ++++++ 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/export/orbax/export/data_processors/jax_data_processor.py b/export/orbax/export/data_processors/jax_data_processor.py index 099e94ec0..35126bd56 100644 --- a/export/orbax/export/data_processors/jax_data_processor.py +++ b/export/orbax/export/data_processors/jax_data_processor.py @@ -22,21 +22,22 @@ import jaxtyping from orbax.export import constants from orbax.export import obm_configs +from orbax.export import typing 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.""" +def _jax_spec_from(spec: Any) -> typing.OrbaxShapeDtypeStruct: + """Converts a ShloTensorSpec to a OrbaxShapeDtypeStruct.""" if isinstance(spec, shlo_function.ShloTensorSpec): if spec.dtype == shlo_function.ShloDType.bf16: - return jax.ShapeDtypeStruct(spec.shape, jax.numpy.bfloat16) - return jax.ShapeDtypeStruct( + return typing.OrbaxShapeDtypeStruct(spec.shape, jax.numpy.bfloat16) + return typing.OrbaxShapeDtypeStruct( spec.shape, shlo_function.shlo_dtype_to_np_dtype(spec.dtype) ) if hasattr(spec, 'shape') and hasattr(spec, 'dtype'): - return jax.ShapeDtypeStruct( + return typing.OrbaxShapeDtypeStruct( shape=tuple(spec.shape), dtype=spec.dtype, ) diff --git a/export/orbax/export/typing.py b/export/orbax/export/typing.py index cf9711af6..02414decf 100644 --- a/export/orbax/export/typing.py +++ b/export/orbax/export/typing.py @@ -54,3 +54,9 @@ class ApplyFnInfo: apply_fn: ApplyFn input_keys: Set[str] output_keys: Set[str] + + +@dataclasses.dataclass(frozen=True, slots=True) +class OrbaxShapeDtypeStruct: + shape: Any + dtype: Any