Rate this Page

★ ★ ★ ★ ★

reference/generated/torchrl.data.vla.VLAAction

Run in Google Colab

VLAAction#

Container for VLA policy outputs.

Parameters:

chunk (torch.Tensor | None) – Continuous action chunk.
tokens (torch.Tensor | None) – Action-token ids in the policy vocabulary.
raw_tokens (torch.Tensor | None) – Backend-native token ids when they differ from action-token window ids.
logits (torch.Tensor | None) – Token logits with a trailing vocabulary dim.
log_probs (torch.Tensor | None) – Log-probabilities of the selected tokens.
mask (torch.Tensor | None) – Optional valid-token/action mask.
padded (bool | None) – Whether variable-length outputs were padded.

Examples

>>> import torch
>>> from torchrl.data.vla import VLAAction
>>> action = VLAAction(chunk=torch.zeros(2, 4, 7), batch_size=[2])
>>> action.chunk.shape
torch.Size([2, 4, 7])

cat(dim: int = 0, *, out=None)#

Concatenates tensordicts along an existing batch dimension dim.

It is equivalent to torch.cat(), and also works under torch.compile(). The tensordicts must have the same keys, or a KeyError is raised, and the same batch size except along dim.

Parameters:

input (sequence of TensorDictBase) – the tensordicts to concatenate.
dim (int, optional) – the batch dimension to concatenate along. Defaults to 0.

Keyword Arguments:

out (TensorDictBase, optional) – a tensordict to write the result into. Defaults to None.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> a = TensorDict(x=torch.zeros(3), batch_size=[3])
>>> b = TensorDict(x=torch.ones(3), batch_size=[3])
>>> TensorDict.cat([a, b]).batch_size
torch.Size([6])

property device: torch.device#

The device of the tensorclass, or None if its leaves may be on different devices.

It is read-only: use to() to move the tensorclass to another device.

classmethod fields() → tuple[dataclasses.Field, ...]#: Returns the fields of the tensorclass.

Deprecated since version 0.15: Use dataclasses.fields(cls) instead. It returns the same tuple of dataclasses.Field objects.

from_any(*, auto_batch_size: bool = False, batch_dims: int | None = None, device: torch.device | None = None, batch_size: torch.Size | None = None)#

Converts a dict, tuple, dataclass, array or file to a tensordict, recursively.

from_any picks the conversion from the type of obj, and applies it to the values inside obj too. It accepts more types than the TensorDict constructor, at the cost of some overhead and of opinionated choices about how each type maps to entries:

a dict or collections.UserDict: from_dict();
a named tuple: from_namedtuple();
a tuple: from_tuple(), with the keys "0", "1", … This includes the results of torch.return_types (max, sort, topk, …), which also emit a FutureWarning. Starting with TensorDict 0.17, the default will change to their field names ("values", "indices", …), as from_namedtuple() gives;
a dataclass instance: from_dataclass() (the result is a tensordict, not a tensorclass). A dataclass type whose fields all have defaults gives a tensordict of the defaults too, with a FutureWarning. Starting with TensorDict 0.17, the default will change to returning the type unchanged, as for any other class;
a NumPy structured array: from_struct_array();
a pandas DataFrame: from_pandas();
an open h5py.File: a PersistentTensorDict (see from_h5()).

A tensor or a tensor collection is returned as it is. A list becomes a tensor if torch.tensor() takes it, and a NonTensorStack otherwise. Any other object is returned unchanged.

Parameters:

obj – the object to convert.

Keyword Arguments:

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, except for a dict, where batch_size wins, and for an h5py.File, where auto_batch_size wins. 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.
device (torch.device, str or int, optional) – the device of the result. Leaves on another device are moved to it. None leaves each leaf where it is, and the result has no device. Defaults to None.
batch_size (torch.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, except for a dict, where batch_size wins and no batch size is computed, and for an h5py.File, where auto_batch_size wins and batch_size is ignored. Defaults to None.

Returns:

A tensordict (a PersistentTensorDict for an HDF5 file), a tensor or a NonTensorStack for a list, or obj itself for the other types.

Examples

>>> import dataclasses
>>> import torch
>>> from tensordict import TensorDict
>>> @dataclasses.dataclass
... class Data:
...     obs: torch.Tensor
...     extra: tuple
>>> data = Data(obs=torch.zeros(3, 4), extra=(torch.ones(3), {"x": torch.ones(3)}))
>>> td = TensorDict.from_any(data, auto_batch_size=True)
>>> td.batch_size
torch.Size([3])
>>> sorted(td.keys(include_nested=True, leaves_only=True), key=str)
[('extra', '0'), ('extra', '1', 'x'), 'obs']
>>> TensorDict.from_any([1, 2, 3])
tensor([1, 2, 3])

from_consolidated()#

Reconstructs a tensordict from a file written by consolidate(filename=...).

The result is consolidated: its leaves are views into one copy-on-write memory map of the file, so in-place writes into them do not reach the file. It is not locked. The batch sizes and names are as they were when the file was written. Same as from_consolidated().

Parameters:: filename (str or Path) – the file to read.

Examples

>>> import tempfile
>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(a=torch.zeros(3), b={"c": torch.ones(3)}, batch_size=[3])
>>> with tempfile.TemporaryDirectory() as path:
...     _ = td.consolidate(filename=f"{path}/td.bin")
...     loaded = TensorDict.from_consolidated(f"{path}/td.bin")
...     print(loaded["b", "c"], loaded.is_consolidated())
tensor([1., 1., 1.]) True

Reads a CSV file into a tensordict, with one entry per column.

Requires pandas or pyarrow: the file is read with pandas.read_csv() when pandas is installed, and with pyarrow.csv.read_csv otherwise. The columns are converted as in from_pandas(). to_csv() writes such a file. Through pandas, a string column is one value for the whole batch (see the note in from_pandas()); through pyarrow, it is a NonTensorStack with one value per row.

Parameters:

path (str or Path) – the path of the CSV file.

Keyword Arguments:

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.
device (torch.device, str or int, optional) – the device of the result. Leaves on another device are moved to it. None leaves each leaf where it is, and the result has no device. Defaults to None.
batch_size (torch.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 gives [num_rows].
separator (str, optional) – if given, splits the column names on it to build nested entries: with separator=".", the column "obs.x" becomes the entry ("obs", "x"). Defaults to None, which keeps the column names as keys.
dtype (torch.dtype, optional) – if given, casts the numeric and boolean columns to it. Defaults to None, which keeps the dtype of each column.
**kwargs – passed to the CSV reader.

Returns:

A tensordict of the class this method is called on (TensorDict for TensorDictBase).

Examples

>>> import os
>>> import tempfile
>>> import torch
>>> from tensordict import TensorDict
>>> data = TensorDict(obs=torch.arange(3.0), reward=torch.ones(3), batch_size=[3])
>>> with tempfile.TemporaryDirectory() as tmpdir:
...     path = os.path.join(tmpdir, "data.csv")
...     data.to_csv(path)
...     td = TensorDict.from_csv(path)
...     td32 = TensorDict.from_csv(path, dtype=torch.float32)
>>> td.batch_size, td["obs"].dtype, td32["obs"].dtype
(torch.Size([3]), torch.float64, torch.float32)

from_dataclass(*, dest_cls: Type | None = None, auto_batch_size: bool = False, batch_dims: int | None = None, as_tensorclass: bool = False, device: torch.device | None = None, batch_size: torch.Size | None = None)#

Converts a dataclass instance to a tensordict, with one entry per field.

Each field value is converted with from_any(), so a nested dataclass or dict becomes a nested tensordict and a number becomes a tensor. With as_tensorclass=True, the conversion is that of the function from_dataclass() instead, which returns a tensorclass.

Parameters:

dataclass – the dataclass instance to convert. With as_tensorclass=True, it can also be a dataclass type. Without it, a dataclass type whose fields all have defaults gives a tensordict of the defaults, with a FutureWarning. Starting with TensorDict 0.17, the default will change to raising a TypeError for it.

Keyword Arguments:

dest_cls (type, optional) – with as_tensorclass=True, the tensorclass to convert to. Ignored otherwise. Defaults to None, which creates a new tensorclass with the fields of the dataclass.
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.
as_tensorclass (bool, optional) – if True, returns a tensorclass instance (or a tensorclass type, for a dataclass type) through from_dataclass(). Defaults to False.
device (torch.device, str or int, optional) – the device of the result. Leaves on another device are moved to it. None leaves each leaf where it is, and the result has no device. Defaults to None.
batch_size (torch.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, an empty batch size.

Returns:

A TensorDict, or a tensorclass instance or type if as_tensorclass=True.

Raises:

TypeError – if dataclass is not a dataclass and as_tensorclass=False, if it is a dataclass type with a field that cannot be read from the class (a field without a default, or a field of a tensorclass type) and as_tensorclass=False, or if both auto_batch_size=True and batch_size are given.

Examples

>>> import dataclasses
>>> import torch
>>> from tensordict import TensorDict
>>> @dataclasses.dataclass
... class Data:
...     obs: torch.Tensor
...     reward: torch.Tensor
>>> data = Data(obs=torch.zeros(3, 4), reward=torch.ones(3))
>>> td = TensorDict.from_dataclass(data, auto_batch_size=True)
>>> type(td).__name__, td.batch_size
('TensorDict', torch.Size([3]))
>>> tc = TensorDict.from_dataclass(data, as_tensorclass=True, batch_size=[3])
>>> type(tc).__name__, tc.batch_size
('Data_tc', torch.Size([3]))

from_h5(*, mode: str = 'r', auto_batch_size: bool = False, batch_dims: int | None = None, batch_size: torch.Size | None = None, non_tensor_as_stack: bool = False)#

Opens an HDF5 file as a PersistentTensorDict.

The data stays in the file: entries are read when they are accessed. Requires h5py. See Save and load tensordicts and PersistentTensorDict.

Parameters:

filename (str or path) – the path to the HDF5 file.

Keyword Arguments:

mode (str, optional) – how the file is opened, as in h5py.File: "r" (read only), "r+" (read and write), and so on. 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 (torch.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 gives each nested tensordict the largest batch size that its entries share, and every nested tensordict batch_size when it is passed. A file written by to_h5() stores the batch sizes and the dim names. Starting with TensorDict 0.17, the default will change to restoring the stored ones, as from_zarr does: with batch_size, the nested tensordicts get their stored batch sizes, and with auto_batch_size=True, the stored names are kept and the batch sizes are inferred from the stored ones, which an empty tensordict or a non-tensor entry keeps. from_h5 warns when this changes its result.
non_tensor_as_stack (bool, optional) – if True, a non-tensor dataset 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 dataset. Defaults to False.

Returns:

A PersistentTensorDict backed by the file.

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, "file.h5")
...     _ = data.to_h5(path)
...     td = TensorDict.from_h5(path)
...     print(td)
PersistentTensorDict(batch_size=[3], fields={key1: f32[], key2: f32[]})

from_json(*, auto_batch_size: bool = False, batch_dims: int | None = None, device: torch.device | None = None, batch_size: torch.Size | None = None, separator: str | None = None, dtype: torch.dtype | None = None, lines: bool = False, **kwargs) → Self#

Reads a JSON file of records into a tensordict, with one entry per field.

The file holds an array of records ([{"a": 1, "b": 2.0}, ...]), or one record per line with lines=True (JSON Lines). Each field becomes an entry with one value per record. The file is read with pandas.read_json() when pandas is installed, and with the json module otherwise. Through pandas, a string field is one value for the whole batch (see the note in from_pandas()); through the json module, it is a NonTensorStack with one value per record. To convert a single JSON object, load it with json.load() and pass it to from_dict(). to_json() writes such a file.

Parameters:

path (str or Path) – the path of the JSON file.

Keyword Arguments:

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.
device (torch.device, str or int, optional) – the device of the result. Leaves on another device are moved to it. None leaves each leaf where it is, and the result has no device. Defaults to None.
batch_size (torch.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 gives [num_rows].
separator (str, optional) – if given, splits the column names on it to build nested entries: with separator=".", the column "obs.x" becomes the entry ("obs", "x"). Defaults to None, which keeps the column names as keys.
dtype (torch.dtype, optional) – if given, casts the numeric and boolean columns to it. Defaults to None, which keeps the dtype of each column.
lines (bool, optional) – if True, reads the file as JSON Lines, one record per line. Defaults to False.
**kwargs – passed to pandas.read_json(). Ignored without pandas.

Returns:

A tensordict of the class this method is called on (TensorDict for TensorDictBase).

Examples

>>> import os
>>> import tempfile
>>> import torch
>>> from tensordict import TensorDict
>>> data = TensorDict(obs=torch.arange(3.0), reward=torch.ones(3), batch_size=[3])
>>> with tempfile.TemporaryDirectory() as tmpdir:
...     path = os.path.join(tmpdir, "data.json")
...     data.to_json(path)
...     td = TensorDict.from_json(path)
...     path_lines = os.path.join(tmpdir, "data.jsonl")
...     data.to_json(path_lines, lines=True)
...     td_lines = TensorDict.from_json(path_lines, lines=True)
>>> td.batch_size, td_lines.batch_size
(torch.Size([3]), torch.Size([3]))

Builds a tensordict by stacking a list of dicts or tensor collections.

Each dict is converted with from_dict() and the results are stacked along a new leading dimension. A tensordict or tensorclass in the list is stacked as it is: the keyword arguments do not apply to it. When the result does not get the batch_size or device asked for, a FutureWarning is emitted: starting with TensorDict 0.17, the default will change to applying them to the result, which raises if the batch size does not fit, as for a list of dicts.

Parameters:

input (list) – the dicts or tensor collections to stack.

Keyword Arguments:

auto_batch_size (bool, optional) – if True, sets the batch size of each element to the largest leading shape that all its entries share. Defaults to None, which behaves like False unless batch_dims is given.
batch_size (torch.Size or list of int, optional) – the batch size of the result, with the length of the list first. Each element gets batch_size[1:]. Defaults to None.
device (torch.device, str or int, optional) – the device of the result. Leaves on another device are moved to it. None leaves each leaf where it is, and the result has no device. Defaults to None.
batch_dims (int, optional) – the maximum number of batch dimensions of the result, the list dimension included. Turns on auto_batch_size when that is None. Exclusive with batch_size. Defaults to None.
names (list of str, optional) – the names of the batch dimensions of the result, one per dimension, with the name of the list dimension first. Defaults to None.
lazy (bool, optional) – if True, the result is a lazy stack; if False, a dense stack. Defaults to None, which builds a dense stack when the elements have the same shape and a lazy stack otherwise. See Stacking and variable-size data.

Raises:

ValueError – if batch_size[0] is not the length of the list.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict.from_list(
...     [{"a": torch.zeros(3)}, {"a": torch.zeros(3)}],
...     batch_size=[2, 3],
...     names=["x", "y"],
... )
>>> td.batch_size
torch.Size([2, 3])
>>> td.names
['x', 'y']

from_module(as_module: bool = False, lock: bool = False, use_state_dict: bool = False, filter_empty: bool = True)#

Returns a tensordict that holds the parameters and buffers of a module.

The entries are the parameters and buffers of the module themselves, not copies. Submodules become nested tensordicts.

Note

This method does not lock the result by default (lock=False). The function tensordict.from_module() locks it by default (lock=True).

Parameters:

module (torch.nn.Module) – the module to get the parameters from.
as_module (bool, optional) – if True, a TensorDictParams instance is returned, which can be used to store parameters within a torch.nn.Module. Defaults to False.
lock (bool, optional) – if True, the resulting tensordict is locked. Defaults to False.
use_state_dict (bool, optional) – if True, the state-dict from the module is used and unflattened into a TensorDict with the tree structure of the model. Use it when the module has state-dict hooks. Defaults to False.
filter_empty (bool, optional) – if True, submodules without parameters or buffers (nn.ReLU, for example) get no entry. If False, they get an empty tensordict. Defaults to True.

Returns:

A TensorDict, or a TensorDictParams if as_module=True.

Examples

>>> from torch import nn
>>> from tensordict import TensorDict
>>> module = nn.Sequential(nn.Linear(3, 4), nn.ReLU())
>>> params = TensorDict.from_module(module)
>>> params["0", "weight"] is module[0].weight
True
>>> list(params.keys())
['0']
>>> params.is_locked
False

from_modules(*, as_module: bool = False, lock: bool = True, use_state_dict: bool = False, lazy_stack: bool = False, expand_identical: bool = False)#

Stacks the parameters and buffers of several modules in one tensordict.

The result has one batch dimension, with one element per module, and the structure that from_module() gives. Use it to run an ensemble of modules that have the same structure with torch.vmap() and to_module() (see the example).

The default, dense result holds new parameters and buffers: the stacks of those of the modules, with the same requires_grad. Pass these new parameters to the optimizer: gradients reach them, not the parameters of the modules. With lazy_stack=True, the result is a lazy stack of the parameters of the modules, which receive the gradients.

Parameters:

*modules (nn.Module) – the modules to get the parameters from. With a dense stack, they must have the same parameters and buffers, with the same shapes.

Keyword Arguments:

as_module (bool, optional) – if True, returns a TensorDictParams, which can be stored as a submodule of a torch.nn.Module. Exclusive with lazy_stack. Defaults to False.
lock (bool, optional) – if True, the result is locked (see Locking and temporary changes). Defaults to True.
use_state_dict (bool, optional) – if True, reads the state_dict() of each module, so that its state-dict hooks run. Defaults to False.
lazy_stack (bool, optional) –
if True, returns a LazyStackedTensorDict of the parameters of the modules; if False, stacks them into new parameters. Exclusive with as_module, and not compatible with lazy modules. Defaults to False.

Warning

A lazy stack stacks the parameters each time an entry is read with get(), which costs memory (N times the size of the parameters, more if a graph is built) and time, and the optimizer holds N times more parameters. Use lazy_stack=True only when you need to train the parameters of the modules themselves.
expand_identical (bool, optional) – if True, a parameter that is the same object in every module is expanded instead of stacked, so that the result shares its memory. Ignored when lazy_stack=True. Defaults to False.

Raises:

RuntimeError – if lazy_stack=True and as_module=True, or if lazy_stack=True and a module has uninitialized (lazy) parameters.

Examples

>>> import torch
>>> from torch import nn
>>> from tensordict import TensorDict
>>> _ = torch.manual_seed(0)
>>> modules = [nn.Linear(3, 4) for _ in range(2)]
>>> params = TensorDict.from_modules(*modules)
>>> params.batch_size, params["weight"].shape
(torch.Size([2]), torch.Size([2, 4, 3]))

Run the two modules on one input with torch.vmap(), through a module on the "meta" device that holds no data. The gradients reach the stacked parameters, not those of the modules:

>>> empty_module = nn.Linear(3, 4, device="meta")
>>> def exec_module(params, x):
...     with params.to_module(empty_module):
...         return empty_module(x)
>>> x = torch.randn(3)
>>> y = torch.vmap(exec_module, (0, None))(params, x)
>>> y.shape
torch.Size([2, 4])
>>> y.sum().backward()
>>> params["weight"].grad is not None, modules[0].weight.grad is None
(True, True)

With lazy_stack=True, the gradients reach the parameters of the modules:

>>> params = TensorDict.from_modules(*modules, lazy_stack=True)
>>> type(params).__name__
'LazyStackedTensorDict'
>>> y = torch.vmap(exec_module, (0, None))(params, x)
>>> y.sum().backward()
>>> modules[0].weight.grad is not None
True

from_namedtuple(*, auto_batch_size: bool = False, batch_dims: int | None = None, device: torch.device | None = None, batch_size: torch.Size | None = None)#

Converts a named tuple to a tensordict whose keys are its field names.

Each field value is converted with from_any(), so a nested named tuple or dict becomes a nested tensordict and a number becomes a tensor. The torch.return_types results of torch.max, torch.sort and the like are accepted too. to_namedtuple() converts back.

Parameters:

named_tuple (NamedTuple) – the named tuple to convert.

Keyword Arguments:

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.
device (torch.device, str or int, optional) – the device of the result. Leaves on another device are moved to it. None leaves each leaf where it is, and the result has no device. Defaults to None.
batch_size (torch.Size or list of int, optional) – the batch size of the result: the leading dimensions that all its entries share. It applies to the nested tensordicts too. Exclusive with auto_batch_size. Defaults to None, an empty batch size.

Returns:

A TensorDict.

Raises:

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

Examples

>>> from tensordict import TensorDict
>>> import torch
>>> data = TensorDict({
...     "a_tensor": torch.zeros((3)),
...     "nested": {"a_tensor": torch.zeros((3)), "a_string": "zero!"}}, [3])
>>> nt = data.to_namedtuple()
>>> nt
GenericDict(a_tensor=tensor([0., 0., 0.]), nested=GenericDict(a_tensor=tensor([0., 0., 0.]), a_string='zero!'))
>>> td = TensorDict.from_namedtuple(nt, auto_batch_size=True)
>>> td.batch_size, td["nested", "a_string"]
(torch.Size([3]), 'zero!')
>>> sorted(TensorDict.from_namedtuple(torch.tensor([2, 0, 1]).max(0)).keys())
['indices', 'values']

Converts a pandas DataFrame to a tensordict, with one entry per column.

The batch size is [num_rows]. Numeric and boolean columns become tensors, a string or object column becomes one NonTensorData that holds the whole column, a categorical column becomes the tensor of its codes, and a datetime column becomes an int64 tensor of nanoseconds. to_pandas() converts back.

Note

A string or object column is one value for the whole batch, so td[i]["name"] gives the values of every row, not the value of row i. To get one value per row, set the column again as a list: td["name"] = df["name"].tolist() stores a NonTensorStack, whose td[i]["name"] is the value of row i. It builds one object per row, about 20 microseconds each, so it is much slower than reading the column for large frames.

Parameters:

dataframe (pd.DataFrame) – the DataFrame to convert.

Keyword Arguments:

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.
device (torch.device, str or int, optional) – the device of the result. Leaves on another device are moved to it. None leaves each leaf where it is, and the result has no device. Defaults to None.
batch_size (torch.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 gives [num_rows].
separator (str, optional) – if given, splits the column names on it to build nested entries: with separator=".", the column "obs.x" becomes the entry ("obs", "x"). Defaults to None, which keeps the column names as keys.
dtype (torch.dtype, optional) – if given, casts the numeric and boolean columns to it. Defaults to None, which keeps the dtype of each column.

Returns:

A tensordict of the class this method is called on (TensorDict for TensorDictBase).

Examples

>>> from tensordict import TensorDict
>>> import pandas as pd
>>> df = pd.DataFrame({"a": [1, 2, 3], "b": [4.0, 5.0, 6.0]})
>>> td = TensorDict.from_pandas(df)
>>> print(td)
TensorDict(batch_size=[3], fields={a: i64[], b: f64[]})

Reads a Parquet file into a tensordict, with one entry per column.

Requires pyarrow or pandas: the file is read with pyarrow.parquet.read_table when pyarrow is installed, and with pandas.read_parquet() otherwise. The columns are converted as in from_pandas(). to_parquet() writes such a file. Through pyarrow, a string column is a NonTensorStack with one value per row; through pandas, it is one value for the whole batch (see the note in from_pandas()).

Parameters:

path (str or Path) – the path of the Parquet file.

Keyword Arguments:

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.
device (torch.device, str or int, optional) – the device of the result. Leaves on another device are moved to it. None leaves each leaf where it is, and the result has no device. Defaults to None.
batch_size (torch.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 gives [num_rows].
separator (str, optional) – if given, splits the column names on it to build nested entries: with separator=".", the column "obs.x" becomes the entry ("obs", "x"). Defaults to None, which keeps the column names as keys.
dtype (torch.dtype, optional) – if given, casts the numeric and boolean columns to it. Defaults to None, which keeps the dtype of each column.
columns (list of str, optional) – the columns to read. Defaults to None, which reads all the columns.
**kwargs – passed to the Parquet reader.

Returns:

A tensordict of the class this method is called on (TensorDict for TensorDictBase).

Examples

>>> import os
>>> import tempfile
>>> import torch
>>> from tensordict import TensorDict
>>> data = TensorDict(obs=torch.arange(3.0), reward=torch.ones(3), batch_size=[3])
>>> with tempfile.TemporaryDirectory() as tmpdir:
...     path = os.path.join(tmpdir, "data.parquet")
...     data.to_parquet(path)
...     td = TensorDict.from_parquet(path)
...     td_obs = TensorDict.from_parquet(path, columns=["obs"])
>>> sorted(td.keys()), sorted(td_obs.keys())
(['obs', 'reward'], ['obs'])

from_pytree(*, batch_size: torch.Size | None = None, auto_batch_size: bool = False, batch_dims: int | None = None)#

Converts a pytree to a tensordict that to_pytree() converts back.

Each list, tuple, named tuple and dict of the pytree becomes a nested tensordict, and its type is stored in a "_pytree_type" entry so that to_pytree() can rebuild the pytree. The elements of a list or tuple get the keys "0", "1", … The subclasses of these types, such as collections.OrderedDict, collections.defaultdict and the torch.return_types results of torch.max and the like, are accepted too. A defaultdict is rebuilt without its default_factory.

A tensor or a tensor collection is returned as it is. Python ints and floats and NumPy arrays become tensors, so to_pytree() returns them as tensors. A dict key that is not a NestedKey (such as an int or an object) is stored in a NonTensorData entry next to its value. Other types raise a NotImplementedError.

Parameters:

pytree – the pytree to convert: a list, tuple, named tuple or dict, nested to any depth, whose leaves are tensors, tensor collections, Python numbers or NumPy arrays.

Keyword Arguments:

batch_size (torch.Size or list of int, optional) – the batch size of the result: the leading dimensions that all its entries share. Defaults to None, which gives an empty batch size unless auto_batch_size=True.
auto_batch_size (bool, optional) – if True, sets the batch size to the largest leading shape that all the entries share. 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.

Note

A module state_dict() round-trips through from_pytree, and its memmap() output loads with load_memmap(path, allow_pickle=False). The state dicts of an optimizer and of a learning-rate scheduler hold None, str values and Python numbers, which from_pytree rejects (None and str) or returns as tensors (numbers): save them with torch.save(sd, path) and load them with torch.load(path, weights_only=True), which round-trips them exactly without unpickling arbitrary objects.

Note

The list fields of a named tuple, and some of its dict and tuple fields, are converted with from_any() and do not round-trip through to_pytree():

a list of numbers comes back as a tensor;
an empty tuple, an empty dict or named tuple, and a dict that holds a str or a list come back as a GenericDict named tuple;
a dict that holds a tuple, a tuple subclass (such as a torch.return_types result), and a tuple that holds anything other than tensors, Python numbers, dicts and named tuples make to_pytree() raise a ValueError.

The other dict, named tuple and tuple fields come back with their own type.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> pytree = {"obs": torch.zeros(3, 4), "pair": (torch.zeros(3), torch.ones(3))}
>>> td = TensorDict.from_pytree(pytree, auto_batch_size=True)
>>> td.batch_size
torch.Size([3])
>>> sorted(td.keys()), sorted(td["pair"].keys())
(['_pytree_type', 'obs', 'pair'], ['0', '1', '_pytree_type'])
>>> type(td.to_pytree()["pair"]).__name__
'tuple'

A pytree with lists, dicts, a named tuple and a dict key that is not a string round-trips:

>>> class Key:
...     pass
>>> key = Key()
>>> pytree = (
...     [torch.arange(3), torch.zeros(2)],
...     {"tensor": torch.ones(2), key: torch.arange(2), "list": [1, 2, 3]},
...     {"named_tuple": TensorDict(two=torch.ones(1)).to_namedtuple()},
... )
>>> pytree_recon = TensorDict.from_pytree(pytree).to_pytree()
>>> def check(v1, v2):
...     assert (v1 == v2).all()
>>> _ = torch.utils._pytree.tree_map(check, pytree, pytree_recon)
>>> key in pytree_recon[1]
True

from_schema(*, batch_size: Sequence[int] | torch.Size | None = None, storage: str | None = None, device=None, **kwargs) → TensorDictBase#

Allocates a zero-filled tensordict from a schema of shapes and dtypes.

Each entry of schema maps a key to an (element_shape, dtype) pair, and the entry gets the shape [*batch_size, *element_shape]. storage selects where the zeros are stored, so that a buffer on disk or on a server can be created without building it in memory first.

Parameters:

schema (dict) – maps each key (a str) to a pair (element_shape, dtype), where element_shape (a list of int or a torch.Size) is the shape of the entry after the batch dimensions.

Keyword Arguments:

batch_size (torch.Size or list of int, optional) – the batch size of the result, prepended to every element shape. Defaults to None, an empty batch size.
storage (str, optional) –
where the entries are stored:
- None: a TensorDict of regular tensors.
- "memmap": memory-mapped tensors, in the directory given by prefix=<dir> (see memmap_like()).
- "h5": a PersistentTensorDict in the HDF5 file given by filename=<path>.
- "zarr": a PersistentTensorDict in the zarr store given by filename=<path or store> (requires zarr>=3.1.1).
- "shared": tensors in shared memory (see share_memory_()).
- "redis" or "dragonfly": a TensorDictStore on a server, through TensorDictStore.from_schema.
Defaults to None.
device (torch.device, str or int, optional) – the device of the result. The "memmap" storage ignores it. Defaults to None.
**kwargs – passed to the backend: prefix and the other arguments of memmap_like() for "memmap", filename for "h5" and "zarr", host, port and the like for "redis" and "dragonfly".

Returns:

A new tensordict of the type that storage selects, filled with zeros.

Raises:

ValueError – if storage is not one of the values above.

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,
...     )
...     print(td_mm.is_memmap())
True

from_struct_array(*, auto_batch_size: bool = False, batch_dims: int | None = None, device: torch.device | None = None, batch_size: torch.Size | None = None) → Self#

Converts a NumPy structured array to a tensordict, with one entry per field.

The batch size is the shape of the array, and each entry holds a field. The conversion is zero-copy: the tensors share memory with the array, so an in-place change to the array changes the tensordict, and the other way around. Two cases copy the data: a device other than the CPU, and a field whose item size does not divide the size of a record, such as the "i8" field of a packed array of "i8" and "f4" fields (12 bytes per record), which torch cannot view. String fields become non-tensor data. to_struct_array() converts back.

Parameters:

struct_array (np.ndarray) – the structured array to convert.

Keyword Arguments:

auto_batch_size (bool, optional) – if True, sets the batch size to the largest leading shape that all the entries share, which can add the dimensions of array fields. 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.
device (torch.device, str or int, optional) – the device of the result. Leaves on another device are moved to it. None leaves each leaf where it is, and the result has no device. Defaults to None.
batch_size (torch.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 uses the shape of the array.

Returns:

A tensordict of the class this method is called on (TensorDict for TensorDictBase).

Raises:

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

Examples

>>> import numpy as np
>>> from tensordict import TensorDict
>>> x = np.array(
...     [("Rex", 9, 81.0), ("Fido", 3, 27.0)],
...     dtype=[("name", "U10"), ("age", "i4"), ("weight", "f4")],
... )
>>> td = TensorDict.from_struct_array(x)
>>> td.batch_size, td["age"]
(torch.Size([2]), tensor([9, 3], dtype=torch.int32))
>>> x["age"] += 1  # the tensordict shares memory with the array
>>> td["age"]
tensor([10,  4], dtype=torch.int32)
>>> bool((td.to_struct_array() == x).all())
True

classmethod from_tensordict(tensordict: TensorDictBase, non_tensordict: dict | None = None, safe: bool | None = None) → Self#

Wraps a tensordict in a new instance of the tensorclass, without copying the leaves.

If tensordict is a TensorDict, the entries of the fields annotated with a tensorclass (or Optional of one) that are tensordicts become instances of that tensorclass, recursively. They are set in a shallow copy of tensordict, which keeps its own entries. An entry with keys that the annotated class does not declare stays a tensordict. Other backends are wrapped as they are: setting a tensorclass in them would write to their storage (e.g. an H5 file).

Parameters:

tensordict (TensorDictBase) – the tensordict that holds the tensor fields.
non_tensordict (dict, optional) – the values of the non-tensor fields. The fields that neither argument holds are set to None. Defaults to None.
safe (bool, optional) – deprecated in 0.15, will be removed in TensorDict 0.17. If False, skips the check that tensordict is a TensorDictBase. By default, the check runs, and another type raises a RuntimeError.

Returns:

A new instance of the class, whose fields are the entries of tensordict.

Examples

>>> import torch
>>> from tensordict import TensorClass
>>> class Pose(TensorClass):
...     q: torch.Tensor
>>> class Obs(TensorClass):
...     a: Pose
>>> obs = Obs(a=Pose(q=torch.zeros(5, 4), batch_size=[5]), batch_size=[5])
>>> type(Obs.from_tensordict(obs.to_tensordict()).a).__name__
'Pose'

from_tuple(*, auto_batch_size: bool = False, batch_dims: int | None = None, device: torch.device | None = None, batch_size: torch.Size | None = None)#

Converts a tuple to a tensordict with the keys "0", "1", …

Each element is converted with from_any(), so a nested tuple or dict becomes a nested tensordict and a number becomes a tensor.

Parameters:

obj (tuple) – the tuple to convert.

Keyword Arguments:

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.
device (torch.device, str or int, optional) – the device of the result. Leaves on another device are moved to it. None leaves each leaf where it is, and the result has no device. Defaults to None.
batch_size (torch.Size or list of int, optional) – the batch size of the result: the leading dimensions that all its entries share. It applies to the nested tensordicts too. Exclusive with auto_batch_size. Defaults to None, an empty batch size.

Returns:

A TensorDict.

Raises:

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

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict.from_tuple((torch.zeros(3, 4), 1), auto_batch_size=True)
>>> td["0"].shape, td["1"], td.batch_size
(torch.Size([3, 4]), tensor(1), torch.Size([]))
>>> td = TensorDict.from_tuple((torch.zeros(3, 4), torch.ones(3)), auto_batch_size=True)
>>> td.batch_size
torch.Size([3])

from_zarr(*, mode: str = 'r', auto_batch_size: bool = False, batch_dims: int | None = None, batch_size: torch.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 (torch.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[]})

fromkeys(value: Any = 0)#

Creates a tensordict with the same value for every key, like dict.fromkeys().

The result is a TensorDict with an empty batch size. A nested key creates nested tensordicts. A Python number becomes a tensor for each key, and a tensor is stored as it is: every key then holds the same tensor object.

Parameters:

keys (list of NestedKey) – the keys of the result.
value (Any, optional) – the value of every key: a tensor, a number, or any value that the TensorDict constructor takes. Defaults to 0.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict.fromkeys(["a", ("b", "c")], torch.zeros(3))
>>> td["b", "c"]
tensor([0., 0., 0.])
>>> td["a"] is td["b", "c"]
True

lazy_stack(dim: int = 0, *, out=None, **kwargs)#

Stacks tensordicts into a lazy stack.

A lazy stack keeps the tensordicts as they are and stacks their values only when you read them. The tensordicts can have different keys and different feature shapes. See Stacking and variable-size data.

Parameters:

input (sequence of TensorDictBase) – the tensordicts to stack.
dim (int, optional) – the new batch dimension. Defaults to 0.

Keyword Arguments:

out (TensorDictBase, optional) – a lazy stack to write the result into. Defaults to None.
**kwargs – passed to LazyStackedTensorDict.lazy_stack, for example stack_dim_name or strict_shape.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> a = TensorDict(x=torch.zeros(3), batch_size=[3])
>>> b = TensorDict(y=torch.ones(3), batch_size=[3])
>>> stack = TensorDict.lazy_stack([a, b])
>>> type(stack).__name__, stack.batch_size
('LazyStackedTensorDict', torch.Size([2, 3]))

load(*args, **kwargs) → Self#

Loads a tensordict that save() or memmap() wrote to disk.

load calls load_memmap(), which documents the other arguments, such as device, mode, subpath and allow_pickle. The leaves of the result read the files lazily. See Save and load tensordicts.

Parameters:

prefix (str or Path) – the directory, or the archive file, to read.
*args – passed to load_memmap().
**kwargs – passed to load_memmap().

Examples

>>> import tempfile
>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(a=torch.arange(3), batch_size=[3])
>>> with tempfile.TemporaryDirectory() as path:
...     _ = td.save(path)
...     loaded = TensorDict.load(path)
...     print(loaded["a"], loaded.batch_size)
MemoryMappedTensor([0, 1, 2]) torch.Size([3])

Loads a tensordict that memmap() or save() wrote, from a directory or an archive.

The leaves of the result read the files lazily: from a directory, each leaf is a MemoryMappedTensor of its file, and from an archive, each leaf is a view into one memory map of the archive. Only the pages of the leaves that you use are read from disk. The result has the class that was saved (a tensorclass is loaded as that tensorclass), and it is not locked. load() is this method. See Save and load tensordicts.

Parameters:

prefix (str or Path) – the directory to read, or an archive file written by save(..., archive=True), by a ".tdz" prefix, or by pack_memmap().
device (torch.device, str or int, optional) – the device of the result. The leaves are copied to it, which reads the files, unless it is the CPU. With "meta", no file is read: the leaves are meta tensors, which show the structure and size of the data. Defaults to None, which keeps the leaves on the CPU, memory-mapped.
non_blocking (bool, optional) – if True, does not synchronize the device after the copies to device. Defaults to False.

Keyword Arguments:

out (TensorDictBase, optional) – a tensordict to load the data into. Its nested tensordicts are reused, but its leaves are replaced by the loaded tensors, so their storage is not reused. Its keys that are not in the saved data are removed. If out has a device, the data is loaded on it, and another device raises a ValueError. To write the data into the storage of out, use out.update_(TensorDict.load_memmap(prefix, device=device)). Defaults to None.
robust_key (bool, optional) – if True, reads the percent-encoded file names that memmap() writes by default, and falls back to the legacy names of keys that are safe file names. If False, reads the legacy layout only. None is the same as True. Defaults to True.
subpath (NestedKey or str, optional) – the nested tensordict to load, as a nested key (("module", "0")) or as a "/"-separated path ("module/0"). Only that subtree is read, from a directory or an archive. A nested key reaches the entries of a saved tensorclass as obj.get(key) does. Defaults to None, which loads the whole tensordict.
mode (str, optional) – how the files are mapped. With "r", the mapping is copy-on-write: in-place writes into the leaves stay in the memory of this process, and the files do not change. For a directory, a pickled copy (a multiprocessing queue, a spawned worker) or a torch.save() of the result reads the files, without these writes: send or save clone() to keep them. With "r+", in-place writes reach the files, and a file that the process cannot write raises a PermissionError. Defaults to None, which maps an archive as with "r", and each file of a directory as with "r+" if the process can write it, and as with "r" otherwise. See the note below.
num_threads (int, optional) – the number of threads that decompress the leaves of a compressed archive. It has no effect on an uncompressed archive or a directory. Defaults to 0, which decompresses them one after the other.
allow_pickle (bool, optional) – whether non-tensor data saved with pickle may be loaded. Pickle can run arbitrary code: pass True only for data you trust, and False for other data, which raises a RuntimeError on pickled data. Defaults to None, which loads pickled data with a FutureWarning. Starting with TensorDict 0.16, the default will change to False. Data saved without pickle loads with any value.

Returns:

The loaded tensordict, or out if it is given.

Raises:

ValueError – if prefix is a file that is not an archive, if subpath names no saved tensordict, or if mode is not None, "r" or "r+".

Note

Load with mode="r" data that you only read: on some network file systems (for example Lustre), page faults on a shared writable mapping take write locks, so readers on different nodes block each other. On Linux, a copy-on-write mapping counts against the memory commit limit, so a file larger than RAM plus swap can fail to map with "r". With "r", memmap_() without a prefix does not write back to prefix.

Note

An archive opened with mode="r+" must have uncompressed, aligned payloads (as written by tensordict without compression) and no nested tensor. In-place writes do not update the CRC-32 checksums of its entries: load_memmap does not check them, but call refresh_archive_checksums() before you give the archive to a tool that does, such as unzip.

Examples

>>> import tempfile
>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(a=torch.zeros(3), nested=TensorDict(b=torch.ones(3)))
>>> with tempfile.TemporaryDirectory() as path:
...     _ = td.memmap(path)
...     loaded = TensorDict.load_memmap(path)
...     print((loaded == td).all(), type(loaded["a"]).__name__)
...     # load one nested tensordict
...     print(TensorDict.load_memmap(path, subpath="nested")["b"])
...     # read the structure without reading the data
...     print(TensorDict.load_memmap(path, device="meta")["a"].device)
True MemoryMappedTensor
MemoryMappedTensor([1., 1., 1.])
meta

With a FakeTensorMode active, the leaves are fake tensors and no file is read:

>>> from torch._subclasses import FakeTensorMode
>>> with tempfile.TemporaryDirectory() as path:
...     _ = td.memmap(path)
...     with FakeTensorMode():
...         print(type(TensorDict.load_memmap(path)["a"]).__name__)
FakeTensor

maybe_dense_stack(dim: int = 0, *, out=None, **kwargs)#

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

The result is a TensorDict if the tensordicts have the same keys and shapes, and a LazyStackedTensorDict otherwise. See Stacking and variable-size data.

Parameters:

input (sequence of TensorDictBase) – the tensordicts to stack.
dim (int, optional) – the new batch dimension. Defaults to 0.

Keyword Arguments:

out (TensorDictBase, optional) – a tensordict to write the result into. Defaults to None.
**kwargs – passed to LazyStackedTensorDict.maybe_dense_stack, for example strict.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> a = TensorDict(x=torch.zeros(3), batch_size=[3])
>>> b = TensorDict(x=torch.ones(3), batch_size=[3])
>>> c = TensorDict(y=torch.ones(3), batch_size=[3])
>>> type(TensorDict.maybe_dense_stack([a, b])).__name__
'TensorDict'
>>> type(TensorDict.maybe_dense_stack([a, c])).__name__
'LazyStackedTensorDict'

stack(dim: int = 0, *, out=None)#

Stacks tensordicts into one dense tensordict along a new dimension dim.

It is equivalent to torch.stack(), and also works under torch.compile(). The tensordicts must have the same batch size and the same keys, or a RuntimeError is raised. To stack tensordicts with different keys or shapes, use lazy_stack(). See Stacking and variable-size data.

Parameters:

input (sequence of TensorDictBase) – the tensordicts to stack.
dim (int, optional) – the new batch dimension. Defaults to 0.

Keyword Arguments:

out (TensorDictBase, optional) – a tensordict to write the result into. Defaults to None.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> a = TensorDict(x=torch.zeros(3), batch_size=[3])
>>> b = TensorDict(x=torch.ones(3), batch_size=[3])
>>> TensorDict.stack([a, b]).batch_size
torch.Size([2, 3])

VLAAction#

Docs

Tutorials

Resources