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, allow_pickle: bool | None = None)#

Opens a zarr store as a PersistentTensorDict.

The data stays in the store: entries are read when they are accessed. Requires zarr>=3.1.1. See Save and load tensordicts and PersistentTensorDict.

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) – how the store is opened: "r" (read only), "r+" (read and write), and so on, as in zarr.open_group(). Defaults to "r".
auto_batch_size (bool, optional) – if True, sets the batch size to the largest leading shape that all the entries share. Exclusive with batch_size. Defaults to False.
batch_dims (int, optional) – with auto_batch_size, the maximum number of batch dimensions. Ignored without auto_batch_size, with a FutureWarning when it would change the result. Starting with TensorDict 0.17, the default will change to turning on auto_batch_size when batch_dims is given, as from_dict does. Defaults to None.
batch_size (Size or list of int, optional) – the batch size of the result: the leading dimensions that all its entries share. Exclusive with auto_batch_size. Defaults to None, which reads the batch size and the names of the batch dimensions that to_zarr() stored, or, for a store written by another tool, gives each nested tensordict the largest batch size that its entries share.
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 dimensions equal the batch size. Otherwise it reads as one NonTensorData that holds the whole array. Defaults to False.
allow_pickle (bool, optional) – whether to decode pickled non-tensor values. False rejects them before decoding. True permits them: use it only for data you trust. None permits them with a FutureWarning when read. Starting with TensorDict 0.17, the default will change to allow_pickle=False. Tensor and JSON entries are unaffected. Defaults to None.

Returns:

A PersistentTensorDict backed by the store.

Raises:

TypeError – if both auto_batch_size=True and batch_size are given.

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(batch_size=[3], fields={key1: f32[], key2: f32[]})