TensorDictBase.from_dataclass#

classmethod TensorDictBase.from_dataclass(dataclass, *, dest_cls: Type | None = None, auto_batch_size: bool = False, batch_dims: int | None = None, as_tensorclass: bool = False, device: device | None = None, batch_size: 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 (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 (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]))

Every tensordict class and every tensorclass has this method. For the other methods, see Methods by task.

TensorDictBase.from_dataclass#

Docs

Tutorials

Resources