Fix serialization case in which `global_shape` instead of `shard_shape` would be chosen as the TensorStore chunk shape in case of degenerate fragments (no primary replica shards on a host).

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

Fix serialization case in which `global_shape` instead of `shard_shape` would be chosen as the TensorStore chunk shape in case of degenerate fragments (no primary replica shards on a host). #1283

Closed copybara-service[bot] closed 3 weeks ago

copybara-service[bot] commented 3 weeks ago

Fix serialization case in which global_shape instead of shard_shape would be chosen as the TensorStore chunk shape in case of degenerate fragments (no primary replica shards on a host).