reference/generated/tensordict.TensorCollection

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorCollection

class tensordict.TensorCollection(*args, **kwargs)

Bases: object

The abstract base class of tensordicts and tensorclasses.

TensorDictBase and TensorClass inherit from it. It cannot be instantiated.

Use it in type hints and isinstance checks for “a tensordict or a tensorclass”. Classes decorated with @tensorclass do not inherit from it, only subclasses of TensorClass do; to check both, use is_tensor_collection().

Examples

>>> import torch
>>> from tensordict import TensorClass, TensorCollection, TensorDict
>>> class Data(TensorClass):
...     x: torch.Tensor
>>> td = TensorDict(x=torch.zeros(3), batch_size=[3])
>>> isinstance(td, TensorCollection)
True
>>> isinstance(Data(x=torch.zeros(3), batch_size=[3]), TensorCollection)
True
>>> isinstance(torch.zeros(3), TensorCollection)
False

Tutorials that use TensorCollection

Using TensorDict for datasets

Using TensorDict for datasets

Functional calls and ensembles with TensorDict

Functional calls and ensembles with TensorDict

Benchmarking TensorDict serialization speed

Benchmarking TensorDict serialization speed

Building tensordicts from streams

Building tensordicts from streams

Using tensorclasses for datasets

Using tensorclasses for datasets

Batched data loading with tensorclasses

Batched data loading with tensorclasses

Simulating particles with a tensorclass

Simulating particles with a tensorclass

Getting, setting, and nesting keys

Getting, setting, and nesting keys

Devices and memory-mapped tensors

Devices and memory-mapped tensors

Building models with TensorDictModule

Building models with TensorDictModule

Filling a TensorDict row by row

Filling a TensorDict row by row

Changing the shape of a TensorDict

Changing the shape of a TensorDict

Slicing, indexing, and masking

Slicing, indexing, and masking

Storing TensorDicts in zarr

Storing TensorDicts in zarr