LazyStackedTensorDict#

LazyStackedTensorDict has the methods of TensorDictBase, listed by task in Methods by task. This page lists only the methods that LazyStackedTensorDict adds, and the overrides that have their own documentation.

class tensordict.LazyStackedTensorDict(*tensordicts: T, stack_dim: int = 0, hook_out: callable | None = None, hook_in: callable | None = None, batch_size: Sequence[int] | None = None, device: device | None = None, names: Sequence[str] | None = None, stack_dim_name: str | None = None, strict_shape: bool = False)#

Bases: TensorDictBase

A stack of tensordicts that keeps each input as a separate tensordict.

torch.stack() and tensordict.stack() return a dense TensorDict. They raise an error when the inputs have different keys, entry shapes or batch sizes. A lazy stack accepts these inputs: it keeps a list of the input tensordicts and stacks their entries only when you read them. Use it for data whose size changes from one sample to the next, such as text or episodes of different lengths. Stacking and variable-size data explains when to use a dense or a lazy stack, and Which container should I use? compares the containers.

Build a lazy stack with tensordict.lazy_stack() or LazyStackedTensorDict(*tds, stack_dim=dim). To get a dense TensorDict when the inputs allow it and a lazy stack otherwise, use tensordict.maybe_dense_stack().

>>> import torch
>>> from tensordict import LazyStackedTensorDict, TensorDict
>>> tds = [TensorDict(a=torch.zeros(3)), TensorDict(a=torch.ones(4))]
>>> td_stack = LazyStackedTensorDict(*tds, stack_dim=0)
>>> td_stack.batch_size
torch.Size([2])
>>> td_stack[0] is tds[0]
True
>>> td_stack.get("a", as_padded_tensor=True)
tensor([[0., 0., 0., 0.],
        [1., 1., 1., 1.]])

A lazy stack behaves as follows:

Indexing along the stack dimension returns the input objects: td_stack[i] is tds[i].
Reading a key (td_stack["a"]) stacks the entries into a new tensor. In-place changes to that tensor do not reach the inputs. Writes through the lazy stack (td_stack["a"] = value, apply_()) do reach them.
If the entry shapes differ, td_stack["a"] raises an error. Pass as_nested_tensor, as_padded_tensor or as_list to get() to read the entries as a nested tensor, a padded tensor or a list. These arguments apply even when the shapes match.
keys() lists only the keys that every input has. The printed lazy stack shows the other keys under exclusive_fields.
The batch sizes of the inputs must have the same number of dimensions. If a dimension differs between inputs, it is -1 in the stacked batch_size. Use strict_shape=True to raise an error instead. An entry shape has -1 in the same way for each dimension that differs between inputs.

Parameters:

*tensordicts (TensorDictBase) – the tensordicts to stack.

Keyword Arguments:

stack_dim (int, optional) – the dimension of the batch size along which the tensordicts are stacked. It must be between -td.ndim - 1 and td.ndim, where td is one of the inputs. Defaults to 0.
hook_out (Callable, optional) –
a callable to execute after get().

Deprecated since version 0.15: hook_out belongs to the vmap() support of lazy stacks and has no replacement.
hook_in (Callable, optional) –
a callable to execute before set().

Deprecated since version 0.15: hook_in belongs to the vmap() support of lazy stacks and has no replacement.
batch_size (Size or list of int, optional) – if no tensordict is given, the batch size of the stacked tensordicts: the empty stack has a batch_size with 0 at stack_dim. If tensordicts are given, it must equal the stacked batch size, or the constructor raises a RuntimeError.
device (device, optional) – the device of an empty stack. If tensordicts are given, the stack takes the device of the first one and ignores this argument. Defaults to None.
names (list of str, optional) – the dimension names of an empty stack. If tensordicts are given, the names come from them, and this argument is ignored with a FutureWarning. To name the dimensions of such a stack, pass stack_dim_name, or set names on the stack. Starting with TensorDict 0.17, the default will change to applying these names to the stack.
stack_dim_name (str, optional) – the name of the stack dimension. Defaults to None.
strict_shape (bool, optional) – if True, the batch sizes of the inputs must be equal, or the constructor raises a RuntimeError. Defaults to False.

Examples

>>> import torch
>>> from tensordict import LazyStackedTensorDict, TensorDict
>>> tds = [TensorDict(a=torch.zeros(3, 4), batch_size=[3]) for _ in range(10)]
>>> td_stack = LazyStackedTensorDict.lazy_stack(tds, -1)
>>> td_stack.shape
torch.Size([3, 10])
>>> td_stack.get("a").shape
torch.Size([3, 10, 4])
>>> td_stack[:, 0] is tds[0]
True

Inputs with different batch sizes give -1 in the batch size:

>>> td_stack = LazyStackedTensorDict(
...     TensorDict(a=torch.zeros(3), batch_size=[3]),
...     TensorDict(a=torch.zeros(4), batch_size=[4]),
... )
>>> td_stack.batch_size
torch.Size([2, -1])

Inputs with different keys keep their own keys:

>>> td_stack = LazyStackedTensorDict(
...     TensorDict(a=torch.zeros(3)), TensorDict(a=torch.zeros(3), b=torch.ones(2))
... )
>>> list(td_stack.keys())
['a']
>>> td_stack[1]["b"]
tensor([1., 1.])

Note

You can assign a list to a key of a lazy stack. The list must have the structure that tensor.tolist() returns: the length of its first level must equal the first dimension of the lazy stack, whether or not it is the stack dimension.

>>> td = LazyStackedTensorDict(TensorDict(), TensorDict(), stack_dim=0)
>>> td["a"] = [torch.ones(2), torch.zeros(1)]
>>> assert td[1]["a"] == torch.zeros(1)
>>> td["b"] = ["a string", "another string"]
>>> assert td[1]["b"] == "another string"

append(tensordict: T) → None#

Appends a tensordict to the lazy stack, like list.append().

Same as insert() at the end of the stack.

Parameters:: tensordict (TensorDictBase) – the tensordict to append. It must have the device of the other members and as many batch dimensions. Its sizes must match theirs, except where the batch size of the lazy stack is -1.

extend(tensordict: list[T] | T) → None#

Appends several tensordicts to the lazy stack, like list.extend().

Parameters:: tensordict (list of TensorDictBase or TensorDictBase) – the tensordicts to append. A tensordict is unbound along the stack dimension and its parts are appended.

classmethod from_dict(input_dict: Dict[str, Dict[NestedKey, Any]] | List[Dict[NestedKey, Any]], *other, auto_batch_size: bool = False, batch_size=None, device=None, batch_dims=None, stack_dim_name=None, stack_dim=0, names: list[str] | None = None)#

Builds a lazy stack from a list of dicts, one tensordict per dict.

Each dict goes through tensordict.TensorDict.from_dict(), then the results are stacked lazily, so the dicts can have different keys and entry shapes. A dict whose keys are "0", "1", … is read as that list.

Parameters:

input_dict (list of dict, or dict of dict) – the data of each stacked tensordict.
*other – passed to tensordict.TensorDict.from_dict().

Keyword Arguments:

auto_batch_size (bool, optional) – if True, sets the batch size of each tensordict to the largest leading shape that all its entries share. Defaults to False.
batch_size (Size or list of int, optional) – the batch size of each stacked tensordict, without the stack dimension. Defaults to None, which gives each tensordict the batch size [] unless its batch size is set automatically.
device (device, str or int, optional) – the device of each stacked tensordict. Defaults to None, which keeps each leaf where it is.
batch_dims (int, optional) – with auto_batch_size, the maximum number of batch dimensions of each tensordict. Defaults to None.
stack_dim_name (str, optional) – the name of the stack dimension. Defaults to None.
stack_dim (int, optional) – the stack dimension. Defaults to 0.
names (list of str, optional) – the names of the dimensions of the result, stack dimension included. Exclusive with stack_dim_name. Without a batch_size, only names[stack_dim] is used today and the members get no names. Starting with TensorDict 0.17, the default will change to giving the others to the members once their batch size is known (a FutureWarning is emitted when this changes the result). Passing other entries whose number differs from the member batch dimensions is deprecated and will be removed in TensorDict 0.17 (a DeprecationWarning is emitted); use stack_dim_name instead.

Returns:

A LazyStackedTensorDict.

Examples

>>> import torch
>>> from tensordict import LazyStackedTensorDict
>>> td = LazyStackedTensorDict.from_dict(
...     [{"a": torch.zeros(3)}, {"a": torch.zeros(4)}], stack_dim_name="t"
... )
>>> td.batch_size, td.names
(torch.Size([2]), ['t'])
>>> td[1]["a"].shape
torch.Size([4])

Builds a lazy stack from a list of dicts, one tensordict per dict.

Same as from_dict() with a list, stacked along dimension 0. The dicts can have different keys and entry shapes.

Parameters:

input (list of dict) – the data of each stacked tensordict.

Keyword Arguments:

auto_batch_size (bool, optional) – if True, sets the batch size of each tensordict to the largest leading shape that all its entries share. Defaults to None, which turns it on when batch_dims is given.
batch_size (Size or list of int, optional) – the batch size of each stacked tensordict, without the stack dimension. Defaults to None, which gives each tensordict the batch size [] unless its batch size is set automatically.
device (device, str or int, optional) – the device of each stacked tensordict. Defaults to None, which keeps each leaf where it is.
batch_dims (int, optional) – the maximum number of batch dimensions of each tensordict when the batch size is set automatically. Defaults to None.
names (list of str, optional) – the names of the dimensions of the result, stack dimension first. See from_dict(). Defaults to None.
lazy (bool, optional) – ignored: the result is always a lazy stack. Defaults to None.

Examples

>>> import torch
>>> from tensordict import LazyStackedTensorDict
>>> td = LazyStackedTensorDict.from_list(
...     [{"a": torch.zeros(3)}, {"a": torch.zeros(4)}], auto_batch_size=True
... )
>>> td.batch_size
torch.Size([2, -1])

get_item_shape(key)#

Returns the shape of an entry of the lazy stack, with -1 where the members differ.

The shape of a tensor entry is read from the members without stacking them. Other entries are read with get(), which stacks them.

Parameters:: key (NestedKey) – the key of the entry.
Raises:: KeyError – if the stack has no entry at key, or only some of its members have one.

Examples

>>> import torch
>>> from tensordict import LazyStackedTensorDict, TensorDict
>>> td = LazyStackedTensorDict(
...     TensorDict(a=torch.zeros(3)), TensorDict(a=torch.zeros(4))
... )
>>> td.get_item_shape("a")
torch.Size([2, -1])

get_nestedtensor(key: NestedKey, default: Any = NO_DEFAULT, *, layout: layout | None = None) → Tensor | TensorCollection#

Returns an entry of the lazy stack as a nested tensor.

Deprecated since version 0.15: Use get(key, as_nested_tensor=True) instead. It differs in three ways: it builds a torch.jagged nested tensor unless you pass layout=torch.strided, whereas this method builds a torch.strided one by default; when a tensor default is given and some tensordicts lack key, it returns default itself, whereas this method uses default as the value of the tensordicts that lack key; and when the stack dim is not 0, it nests the tensors along the stack dim and puts that dim first, whereas this method raises.

Parameters:

key (NestedKey) – the entry to nest.
default (Any, optional) – the value for the tensordicts that lack key. If it is a tensor, it is nested with the entries of the other tensordicts. Otherwise, it is returned as is. Without a default, a missing key raises a KeyError.

Keyword Arguments:

layout (layout, optional) – the layout of the nested tensor. Defaults to None, which builds a torch.strided one.

Examples

>>> import torch
>>> from tensordict import LazyStackedTensorDict, TensorDict
>>> td0 = TensorDict({"a": torch.zeros(4), "b": torch.zeros(4)}, [])
>>> td1 = TensorDict({"a": torch.ones(5)}, [])
>>> td = LazyStackedTensorDict.lazy_stack([td0, td1], 0)
>>> a = td.get_nestedtensor("a")
>>> # using a tensor as default uses this default to build the nested tensor
>>> b = td.get_nestedtensor("b", default=torch.ones(4))
>>> assert (b[1] == 1).all()
>>> # using anything else as default returns the default
>>> b2 = td.get_nestedtensor("b", None)
>>> assert b2 is None

property hook_in: Callable | None#: The callable that the lazy stack applies to the values it writes.

Deprecated since version 0.15: hook_in belongs to the vmap() support of lazy stacks and has no replacement.

property hook_out: Callable | None#: The callable that the lazy stack applies to the values it returns.

Deprecated since version 0.15: hook_out belongs to the vmap() support of lazy stacks and has no replacement.

insert(index: int, tensordict: T) → None#

Inserts a tensordict into the lazy stack at index, like list.insert().

The lazy stack is changed in place and its batch size grows by one along the stack dimension. The new tensordict is stored as is, not copied.

Parameters:

index (int) – the position of the new tensordict along the stack dimension.
tensordict (TensorDictBase) – the tensordict to insert. It must have the device of the other members and as many batch dimensions. Its sizes must match theirs, except where the batch size of the lazy stack is -1.

Raises:

TypeError – if tensordict is not a tensordict.
ValueError – if the device of tensordict differs from that of the members, or if its batch size does not fit the lazy stack.
RuntimeError – if the lazy stack is locked.

Examples

>>> import torch
>>> from tensordict import LazyStackedTensorDict, TensorDict
>>> td = LazyStackedTensorDict(TensorDict(a=torch.zeros(3)))
>>> td.insert(0, TensorDict(a=torch.ones(3)))
>>> td.batch_size
torch.Size([2])
>>> td[0]["a"]
tensor([1., 1., 1.])

classmethod lazy_stack(items: Sequence[TensorDictBase], dim: int = 0, *, device: device | str | int | None = None, out: T | None = None, stack_dim_name: str | None = None, strict_shape: bool = False, as_list: bool = False, as_padded_tensor: bool = False, as_nested_tensor: bool = False, padding_side: str = 'right', layout: layout | None = None, padding_value: float | int | bool = 0.0) → Self#

Stacks tensordicts into a lazy stack.

The result keeps each input as a separate tensordict, so the inputs can have different keys, entry shapes and batch sizes. See LazyStackedTensorDict for the behavior of the result, and Stacking and variable-size data for the choice between a dense and a lazy stack.

If items holds tensorclass instances of one class, the result is an instance of that class that holds a lazy stack. If items holds tensors, they are stacked with torch.stack(), unless one of as_list, as_padded_tensor or as_nested_tensor is set. These arguments, layout, padding_side and padding_value apply only to a list of tensors. To read padded or nested entries from a lazy stack of tensordicts, pass them to get() instead.

Parameters:

items (Sequence of TensorDictBase) – the tensordicts to stack. It must not be empty, or a RuntimeError is raised.
dim (int, optional) – the dimension along which to stack. Defaults to 0. A dim below -batch_dims - 1 warns and wraps around once more (and raises if it is still out of range). Starting with TensorDict 0.17, the default will change to raising an IndexError, as for a tensor.

Keyword Arguments:

device (device, optional) – not used: the stack takes the device of the first item.
out (TensorDictBase, optional) – a lazy stack to write the data into. Its batch size must be the stacked batch size. Defaults to None.
stack_dim_name (str, optional) – the name of the stack dimension. Defaults to None.
strict_shape (bool, optional) – if True, the batch sizes of the items must be equal, or a RuntimeError is raised. Defaults to False.
as_list (bool, optional) – tensors only. If True, the tensors are returned as a list. Exclusive with as_padded_tensor and as_nested_tensor. Defaults to False.
as_padded_tensor (bool, optional) – tensors only. If True, the tensors are padded to the largest shape and stacked into one tensor: each dim is padded to its largest size. padding_value sets the padding value and padding_side sets the side of every padded dim. Exclusive with as_list and as_nested_tensor. Defaults to False.
as_nested_tensor (bool, optional) – tensors only. If True, the tensors are returned as a nested tensor (torch.nested). layout sets its layout. Tensors that are nested tensors themselves raise a ValueError, as a nested tensor has a single nested level. Exclusive with as_list and as_padded_tensor. Defaults to False.
layout (layout, optional) – the layout of the nested tensor when as_nested_tensor=True: torch.jagged or torch.strided. Defaults to torch.jagged.
padding_side (str, optional) – the side to pad with as_padded_tensor: "left" or "right". Defaults to "right".
padding_value (float, int or bool, optional) – the padding value with as_padded_tensor. Defaults to 0.0.

Returns:

A LazyStackedTensorDict, or a tensorclass instance if items holds tensorclass instances.

Examples

>>> import torch
>>> from tensordict import LazyStackedTensorDict, TensorDict
>>> td_stack = LazyStackedTensorDict.lazy_stack(
...     [TensorDict(a=torch.zeros(3)), TensorDict(a=torch.zeros(4))]
... )
>>> type(td_stack).__name__
'LazyStackedTensorDict'
>>> td_stack.get("a", as_list=True)
[tensor([0., 0., 0.]), tensor([0., 0., 0., 0.])]

classmethod maybe_dense_stack(items: Sequence[TensorDictBase], dim: int = 0, out: T | None = None, strict: bool = False) → Self#

Stacks tensordicts densely if possible, or into a lazy stack otherwise.

Use this function when the inputs usually match, but can differ. If every item has the same keys and the same entry shapes, the result is the dense stack that torch.stack() returns. Otherwise, the result is the lazy stack that lazy_stack() returns. See Stacking and variable-size data.

Parameters:

items (Sequence of TensorDictBase or tensorclass instances) – the tensordicts to stack.
dim (int, optional) – the new dimension. Defaults to 0. A dim outside [-batch_dims - 1, batch_dims] warns, and is used on the dims of the leaves. Starting with TensorDict 0.17, the default will change to raising an IndexError, as for a tensor.
out (TensorDictBase, optional) – a tensordict to write the result into. Defaults to None.
strict (bool, optional) – used only with out. If True, out must have the same keys as the items, or a RuntimeError is raised. Defaults to False.

Returns:

A TensorDict (or the tensorclass of the items) if the items can be stacked densely, a LazyStackedTensorDict otherwise.

Examples

>>> from tensordict import LazyStackedTensorDict, TensorDict
>>> td0 = TensorDict({"a": 0}, [])
>>> td1 = TensorDict({"b": 0}, [])
>>> out = LazyStackedTensorDict.maybe_dense_stack([td0, td0])  # same keys: dense stack
>>> assert type(out) is TensorDict and out.shape == (2,)
>>> out = LazyStackedTensorDict.maybe_dense_stack([td0, td1])  # different keys: lazy stack
>>> assert isinstance(out, LazyStackedTensorDict) and out.shape == (2,)

non_tensor_items(include_nested: bool = False)#

Returns the non-tensor leaves as (key, value) pairs, each stacked across the members.

Only the keys that every member holds are listed. If include_nested is True, the nested tensordicts are searched as well.

valid_keys(include_nested: bool = False, leaves_only: bool = False, is_leaf: Callable[[Type], bool] | None = None, *, sort: bool = False) → _LazyStackedTensorDictKeysView#: Returns the keys of the lazy stack, like keys().

Deprecated since version 0.15: Use keys() instead, which takes the same arguments. Will be removed in TensorDict 0.17.

Tutorials that use `LazyStackedTensorDict`#

Building tensordicts from streams

Changing the shape of a TensorDict

LazyStackedTensorDict#

Tutorials that use `LazyStackedTensorDict`#

Docs

Tutorials

Resources

LazyStackedTensorDict#

Tutorials that use LazyStackedTensorDict#

Docs

Tutorials

Resources

Tutorials that use `LazyStackedTensorDict`#