TensorDictBase.from_schema#

classmethod TensorDictBase.from_schema(schema: dict[str, tuple[list[int] | Size, dtype]], *, batch_size: Sequence[int] | Size | None = None, storage: str | None = None, device=None, **kwargs) → TensorDictBase#

Pre-allocate a zero-filled TensorDict from a schema.

Creates a TensorDictBase whose storage backend is selected by storage. Each entry in schema maps a field name to an (element_shape, dtype) pair; the full stored shape is [*batch_size, *element_shape].

Parameters:

schema – Mapping from field name to (element_shape, dtype). element_shape is the per-element shape (excluding batch_size).

Keyword Arguments:

batch_size – Overall batch dimensions prepended to every element shape. Defaults to ().
storage (str or None) –
Backend selector:
- None – plain TensorDict with regular tensors.
- "memmap" – memory-mapped tensors on disk. Pass prefix=<dir> in kwargs.
- "h5" – HDF5 via PersistentTensorDict. Pass filename=<path> in kwargs.
- "zarr" – zarr (requires zarr>=3.0) via PersistentTensorDict. Pass filename=<path or store> in kwargs.
- "shared" – CPU shared-memory tensors.
- "redis" / "dragonfly" – delegates to TensorDictStore.from_schema().
device – Device for the resulting tensors (ignored by some backends).
**kwargs – Backend-specific arguments forwarded to the underlying constructor (e.g. prefix for memmap, filename for h5, host/port for redis).

Returns:

A new TensorDictBase subclass instance with pre-allocated (zero-filled) keys.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict.from_schema(
...     {"obs": ([84, 84, 3], torch.uint8),
...      "reward": ([], torch.float32)},
...     batch_size=[1000],
... )
>>> td["obs"].shape
torch.Size([1000, 84, 84, 3])

>>> import tempfile
>>> with tempfile.TemporaryDirectory() as d:
...     td_mm = TensorDict.from_schema(
...         {"obs": ([4], torch.float32)},
...         batch_size=[8],
...         storage="memmap",
...         prefix=d,
...     )
...     assert td_mm.is_memmap()