Refactor `build_array_tspec_for_write` into a dataclass.

google / orbax

Orbax provides common checkpointing and persistence utilities for JAX users

https://orbax.readthedocs.io/

Apache License 2.0

305 stars 36 forks source link

Closed copybara-service[bot] closed 1 month ago

copybara-service[bot] commented 1 month ago

Refactor build_array_tspec_for_write into a dataclass.

This allows to expose some properties of the spec easier (if ocdbt is used, Zarr version, chunk shape).