tensordict package

The TensorDict class simplifies the process of passing multiple tensors from module to module by packing them in a dictionary-like object that inherits features from regular pytorch tensors.

TensorDictBase()	TensorDictBase is an abstract parent class for TensorDicts, a torch.Tensor data container.
TensorDict([source, batch_size, device, ...])	A batched dictionary of tensors.
LazyStackedTensorDict(*tensordicts[, ...])	A Lazy stack of TensorDicts.
PersistentTensorDict(*[, batch_size, ...])	Persistent TensorDict implementation.
TensorDictParams([parameters, no_convert, lock])	A Wrapper for TensorDictBase with Parameter Exposure.
get_defaults_to_none([set_to_none])	Returns the status of get default value.

Constructors and handlers

The library offers a few method to interact with other data structures such as numpy structured arrays, namedtuples or h5 files. The library also exposes dedicated functions to manipulate tensordicts such as save, load, stack or cat.

cat(input[, dim, out])	Concatenates tensordicts into a single tensordict along the given dimension.
default_is_leaf(cls)	Returns True if a type is not a tensor collection (tensordict or tensorclass).
from_any(obj, *[, auto_batch_size, ...])	Converts any object to a TensorDict.
from_consolidated(filename)	Reconstructs a tensordict from a consolidated file.
from_dict(d, *[, auto_batch_size, ...])	Converts a dictionary to a TensorDict.
from_h5(h5_file, *[, auto_batch_size, ...])	Converts an HDF5 file to a TensorDict.
from_module(module[, as_module, lock, ...])	Copies the params and buffers of a module in a tensordict.
from_modules(*modules[, as_module, lock, ...])	Retrieves the parameters of several modules for ensebmle learning/feature of expects applications through vmap.
from_namedtuple(named_tuple, *[, ...])	Converts a namedtuple to a TensorDict.
from_pytree(pytree, *[, batch_size, ...])	Converts a pytree to a TensorDict instance.
from_struct_array(struct_array, *[, ...])	Converts a structured numpy array to a TensorDict.
from_tuple(obj, *[, auto_batch_size, ...])	Converts a tuple to a TensorDict.
fromkeys(keys[, value])	Creates a tensordict from a list of keys and a single value.
is_batchedtensor(arg0)
is_leaf_nontensor(cls)	Returns True if a type is not a tensor collection (tensordict or tensorclass) or is a non-tensor.
lazy_stack(input[, dim, out])	Creates a lazy stack of tensordicts.
load(prefix[, device, non_blocking, out])	Loads a tensordict from disk.
load_memmap(prefix[, device, non_blocking, out])	Loads a memory-mapped tensordict from disk.
maybe_dense_stack(input[, dim, out])	Attempts to make a dense stack of tensordicts, and falls back on lazy stack when required..
memmap(data[, prefix, copy_existing, ...])	Writes all tensors onto a corresponding memory-mapped Tensor in a new tensordict.
save(data[, prefix, copy_existing, ...])	Saves the tensordict to disk.
stack(input[, dim, out])	Stacks tensordicts into a single tensordict along the given dimension.

TensorDict as a context manager

TensorDict can be used as a context manager in situations where an action has to be done and then undone. This include temporarily locking/unlocking a tensordict

>>> data.lock_()  # data.set will result in an exception
>>> with data.unlock_():
...     data.set("key", value)
>>> assert data.is_locked()

or to execute functional calls with a TensorDict instance containing the parameters and buffers of a model:

>>> params = TensorDict.from_module(module).clone()
>>> params.zero_()
>>> with params.to_module(module):
...     y = module(x)

In the first example, we can modify the tensordict data because we have temporarily unlocked it. In the second example, we populate the module with the parameters and buffers contained in the params tensordict instance, and reset the original parameters after this call is completed.

Memory-mapped tensors

tensordict offers the MemoryMappedTensor primitive which allows you to work with tensors stored in physical memory in a handy way. The main advantages of MemoryMappedTensor are its ease of construction (no need to handle the storage of a tensor), the possibility to work with big contiguous data that would not fit in memory, an efficient (de)serialization across processes and efficient indexing of stored tensors.

If all workers have access to the same storage (both in multiprocess and distributed settings), passing a MemoryMappedTensor will just consist in passing a reference to a file on disk plus a bunch of extra meta-data for reconstructing it. The same goes with indexed memory-mapped tensors as long as the data-pointer of their storage is the same as the original one.

Indexing memory-mapped tensors is much faster than loading several independent files from the disk and does not require to load the full content of the array in memory. However, physical storage of PyTorch tensors should not be any different:

>>> my_images = MemoryMappedTensor.empty((1_000_000, 3, 480, 480), dtype=torch.unint8)
>>> mini_batch = my_images[:10]  # just reads the first 10 images of the dataset

MemoryMappedTensor(source, *[, dtype, ...])

A Memory-mapped Tensor.

Pointwise Operations

Tensordict supports various pointwise operations, allowing you to perform element-wise computations on the tensors stored within it. These operations are similar to those performed on regular PyTorch tensors.

Supported Operations

The following pointwise operations are currently supported:

• Left and right addition (+)

• Left and right subtraction (-)

• Left and right multiplication (*)

• Left and right division (/)

• Left power (**)

Many other ops, like clamp(), sqrt() etc. are supported.

Performing Pointwise Operations

You can perform pointwise operations between two Tensordicts or between a Tensordict and a tensor/scalar value.

Example 1: Tensordict-Tensordict Operation

>>> import torch
>>> from tensordict import TensorDict
>>> td1 = TensorDict(
...     a=torch.randn(3, 4),
...     b=torch.zeros(3, 4, 5),
...     c=torch.ones(3, 4, 5, 6),
...     batch_size=(3, 4),
... )
>>> td2 = TensorDict(
...     a=torch.randn(3, 4),
...     b=torch.zeros(3, 4, 5),
...     c=torch.ones(3, 4, 5, 6),
...     batch_size=(3, 4),
... )
>>> result = td1 * td2

In this example, the * operator is applied element-wise to the corresponding tensors in td1 and td2.

Example 2: Tensordict-Tensor Operation

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(
...     a=torch.randn(3, 4),
...     b=torch.zeros(3, 4, 5),
...     c=torch.ones(3, 4, 5, 6),
...     batch_size=(3, 4),
... )
>>> tensor = torch.randn(4)
>>> result = td * tensor

ere, the * operator is applied element-wise to each tensor in td and the provided tensor. The tensor is broadcasted to match the shape of each tensor in the Tensordict.

Example 3: Tensordict-Scalar Operation

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(
...     a=torch.randn(3, 4),
...     b=torch.zeros(3, 4, 5),
...     c=torch.ones(3, 4, 5, 6),
...     batch_size=(3, 4),
... )
>>> scalar = 2.0
>>> result = td * scalar

In this case, the * operator is applied element-wise to each tensor in td and the provided scalar.

Broadcasting Rules

When performing pointwise operations between a Tensordict and a tensor/scalar, the tensor/scalar is broadcasted to match the shape of each tensor in the Tensordict: the tensor is broadcast on the left to match the tensordict shape, then individually broadcast on the right to match the tensors shapes. This follows the standard broadcasting rules used in PyTorch if one thinks of the TensorDict as a single tensor instance.

For example, if you have a Tensordict with tensors of shape (3, 4) and you multiply it by a tensor of shape (4,), the tensor will be broadcasted to shape (3, 4) before the operation is applied. If the tensordict contains a tensor of shape (3, 4, 5), the tensor used for the multiplication will be broadcast to (3, 4, 5) on the right for that multiplication.

If the pointwise operation is executed across multiple tensordicts and their batch-size differ, they will be broadcasted to a common shape.

Efficiency of pointwise operations

When possible, torch._foreach_<op> fused kernels will be used to speed up the computation of the pointwise operation.

Handling Missing Entries

When performing pointwise operations between two Tensordicts, they must have the same keys. Some operations, like add(), have a default keyword argument that can be used to operate with tensordict with exclusive entries. If default=None (the default), the two Tensordicts must have exactly matching key sets. If default="intersection", only the intersecting key sets will be considered, and other keys will be ignored. In all other cases, default will be used for all missing entries on both sides of the operation.

Utils

utils.expand_as_right(tensor, dest)	Expand a tensor on the right to match another tensor shape.
utils.expand_right(tensor, shape)	Expand a tensor on the right to match a desired shape.
utils.isin(input, reference, key[, dim])	Tests if each element of key in input dim is also present in the reference.
utils.remove_duplicates(input, key[, dim, ...])	Removes indices duplicated in key along the specified dimension.
is_batchedtensor(arg0)
is_tensor_collection(datatype)	Checks if a data object or a type is a tensor container from the tensordict lib.
make_tensordict([input_dict, batch_size, ...])	Returns a TensorDict created from the keyword arguments or an input dictionary.
merge_tensordicts(*tensordicts[, callback_exist])	Merges tensordicts together.
pad(tensordict, pad_size[, value])	Pads all tensors in a tensordict along the batch dimensions with a constant value, returning a new tensordict.
pad_sequence(list_of_tensordicts[, pad_dim, ...])	Pads a list of tensordicts in order for them to be stacked together in a contiguous format.
dense_stack_tds(td_list[, dim])	Densely stack a list of TensorDictBase objects (or a LazyStackedTensorDict) given that they have the same structure.
set_lazy_legacy(mode)	Sets the behaviour of some methods to a lazy transform.
lazy_legacy([allow_none])	Returns True if lazy representations will be used for selected methods.
parse_tensor_dict_string(s)	Parse a TensorDict repr to a TensorDict.