Reading a slice of an arrays across PyTrees

[darsnack]

I have a checkpoint which is a single PyTree. Each leaf has a leading batch dimension. I want to be able to restore the i-th slice of this leading dimension only. The end result should be a PyTree of the same structure as the checkpoint, but each leaf is the i-th slice. Is there a way to write a custom handler to facilitate this? I know that Zarr3 let's you read only a chunk of the data from disk.

[niketkumar]

Please consider the following options:

• PyTreeRestoreArgs.transforms in PyTreeCheckpointHandler.restore.

• Write a custom TypeHandler to handle ndarray or jax.Array. See https://orbax.readthedocs.io/en/latest/custom_handlers.html#typehandler for details.