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) – 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.
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 dims 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.randn(3, 4)}, batch_size=[3])
...     for _ in range(10)]
>>> td_stack = LazyStackedTensorDict.lazy_stack(tds, -1)
>>> print(td_stack.shape)
torch.Size([3, 10])
>>> print(td_stack.get("a").shape)
torch.Size([3, 10, 4])
>>> print(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#

Append a TensorDict onto the stack.

Analogous to list.append. The appended TensorDict must have compatible batch_size and device. The append operation is in-place, nothing is returned.

Parameters:: tensordict (TensorDictBase) – The TensorDict to be appended onto the stack.

extend(tensordict: list[T] | T) → None#: Extends the lazy stack with new tensordicts.

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 dicts, or dict of dicts) – the data of each stacked tensordict.
*other – passed to tensordict.TensorDict.from_dict().

Keyword Arguments:

auto_batch_size (bool, optional) – if True, the batch size of each tensordict is computed from its data. Defaults to False.
batch_size (Size, optional) – the batch size of each stacked tensordict.
device (device, optional) – the device of each stacked tensordict.
batch_dims (int, optional) – with auto_batch_size, the number of batch dimensions to keep.
stack_dim_name (str, optional) – the name of the stack dimension.
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.

Same as from_dict() with a list. lazy is ignored: the result is always a lazy stack. See tensordict.TensorDictBase.from_list() for the other arguments.

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)#

Gets the shape of an item in the lazy stack.

Heterogeneous dimensions are returned as -1.

This implementation is inefficient as it will attempt to stack the items to compute their shape, and should only be used for printing.

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

Returns a nested tensor when stacking cannot be achieved.

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 default value to return in case the key isn’t in all sub-tensordicts.

Note

In case the default is a tensor, this method will attempt the construction of a nestedtensor with it. Otherwise, the default value will be returned.

Keyword Arguments:

layout (layout, optional) – the layout for the nested tensor.

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#

Insert a TensorDict into the stack at the specified index.

Analogous to list.insert. The inserted TensorDict must have compatible batch_size and device. Insertion is in-place, nothing is returned.

Parameters:

index (int) – The index at which the new TensorDict should be inserted.
tensordict (TensorDictBase) – The TensorDict to be inserted into the stack.

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 in a LazyStackedTensorDict.

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.

Parameters:

items (Sequence of TensorDictBase instances) – the tensordicts to stack. It must not be empty, or a RuntimeError is raised.
dim (int, optional) – the dim along which to perform the lazy 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 LazyStackedTensorDict where to write the data. Its batch size must be the stacked batch size.
stack_dim_name (str, optional) – a name for the stacked dimension.
strict_shape (bool, optional) – if True, the batch sizes of the items must be equal, or a RuntimeError is raised. Defaults to False.

The following keyword arguments apply only when items is a list of tensors. With tensordicts, they have no effect. To read padded or nested entries from a lazy stack of tensordicts, pass them to get() instead.

Keyword Arguments:

as_list (bool, optional) – 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) – 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) – 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) – The side of padding. Must be “left” or “right”. Defaults to “right”.
padding_value (scalar or bool, optional) – The padding value. 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 onto a LazyStackedTensorDict 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.
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,)

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.

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`#