Skip to content

Commit

Permalink
Fix orbax error after new release
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 699508213
Conchylicultor authored and The etils Authors committed Nov 23, 2024

Unverified

This commit is not signed, but one or more authors requires that any commit attributed to them is signed.
1 parent 8a5685a commit c412811
Showing 2 changed files with 4 additions and 3 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -10,6 +10,7 @@ Changelog follow https://keepachangelog.com/ format.

* `enp`:
* Make `enp.testing.parametrize_xnp()` import only requested xnp modules.
* Fix orbax error when inspecting specs of an orbax checkpoint.
* `ecolab`:
* `ecolab.inspect`: Proto are better displayed (hide attributes
`DESCRIPTOR`, `Extensions` in sub-section)
6 changes: 3 additions & 3 deletions etils/enp/array_spec.py
Original file line number Diff line number Diff line change
@@ -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 import checkpoint as ocp # pylint: disable=g-import-not-at-top # pytype: disable=import-error
from orbax.checkpoint._src.serialization import type_handlers # pylint: disable=g-import-not-at-top # pytype: disable=import-error

return isinstance(
array,
(
ocp.type_handlers.ArrayMetadata,
ocp.type_handlers.ScalarMetadata,
type_handlers.ArrayMetadata,
type_handlers.ScalarMetadata,
),
)

0 comments on commit c412811

Please sign in to comment.