Fix orbax error in open-source

PiperOrigin-RevId: 735379386

Author: Conchylicultor

CHANGELOG.md

@@ -8,6 +8,11 @@ Changelog follow https://keepachangelog.com/ format.
 
 ## [Unreleased]
 
+## [1.12.2] - 2025-03-10
+
+*   `enp`:
+    *   Fix orbax ArraySpec issue.
+
 ## [1.12.1] - 2025-02-07
 
 *   `enp`:
@@ -486,7 +491,8 @@ Changelog follow https://keepachangelog.com/ format.
 
 <!-- mdlint off(LINK_UNUSED_ID) -->
 
-[Unreleased]: https://github.com/google/etils/compare/v1.12.1...HEAD
+[Unreleased]: https://github.com/google/etils/compare/v1.12.2...HEAD
+[1.12.2]: https://github.com/google/etils/compare/v1.12.1...v1.12.2
 [1.12.1]: https://github.com/google/etils/compare/v1.12.0...v1.12.1
 [1.12.0]: https://github.com/google/etils/compare/v1.11.0...v1.12.0
 [1.11.0]: https://github.com/google/etils/compare/v1.10.0...v1.11.0

etils/__init__.py

@@ -16,7 +16,7 @@
 
 # A new PyPI release will be pushed everytime `__version__` is increased
 # When changing this, also update the CHANGELOG.md
-__version__ = '1.12.1'
+__version__ = '1.12.2'
 
 # Do NOT add anything to this file. This is left empty on purpose.
 # Users should import subprojects directly.

etils/enp/array_spec.py

@@ -182,12 +182,12 @@ def _is_pygrain(array: Array) -> bool:
 def _is_orbax(array: Array) -> bool:
   if 'orbax.checkpoint' not in sys.modules:
     return False
-  from orbax.checkpoint._src.serialization import type_handlers  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
+  from orbax.checkpoint.metadata import value  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
 
   return isinstance(
       array,
       (
-          type_handlers.ArrayMetadata,
-          type_handlers.ScalarMetadata,
+          value.ArrayMetadata,
+          value.ScalarMetadata,
       ),
   )