TensorDictBase.from_zarr#

classmethod TensorDictBase.from_zarr(filename, *, mode: str = 'r', auto_batch_size: bool = False, batch_dims: int | None = None, batch_size: Size | None = None, non_tensor_as_stack: bool = False)#

Creates a PersistentTensorDict from a zarr store.

Requires zarr>=3.0 to be installed.

Parameters:

filename (str, path or zarr store) – The path to the zarr store (a directory), or a zarr.abc.store.Store instance (e.g. a zarr.storage.ZipStore).

Keyword Arguments:

mode (str, optional) – Reading mode. Defaults to "r".
auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.
batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None (full batch-size at each level).
batch_size (Size, optional) – The batch size of the TensorDict. Defaults to None (read from the metadata written by to_zarr(), or automatically determined for stores written by other tools).
non_tensor_as_stack (bool, optional) – if True, a non-tensor array that tensordict did not write, such as a string column with one string per row, reads as a NonTensorStack of its rows when its leading dims equal the batch size. Otherwise it reads as one NonTensorData that holds the whole array. Defaults to False.

Returns:

A PersistentTensorDict representation of the input zarr store.

Examples

>>> import os
>>> import tempfile
>>> import torch
>>> from tensordict import TensorDict
>>> data = TensorDict(key1=torch.zeros(3), key2=torch.ones(3), batch_size=[3])
>>> with tempfile.TemporaryDirectory() as tmpdir:
...     path = os.path.join(tmpdir, "store.zarr")
...     _ = data.to_zarr(path)
...     td = TensorDict.from_zarr(path)
...     print(td)
PersistentTensorDict(
    fields={
        key1: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        key2: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)