reference/generated/tensordict.TensorDictBase.reshape

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.reshape

TensorDictBase.reshape(*args, **kwargs) → Self

Returns a tensordict whose batch dimensions are reshaped to shape.

Each leaf is reshaped with torch.Tensor.reshape(), and keeps its feature dimensions. As with torch.Tensor.reshape(), a leaf of the result is a view of the original leaf when the memory layout allows it, and a copy otherwise. The result has no names.

Parameters:

*shape (int) – the new batch size, as integers or as one list, tuple or torch.Size. One dimension can be -1. The shape is positional only: reshape(shape=...) raises TypeError.

Keyword Arguments:

inplace (bool, optional) – if True, the tensordict keeps its identity and keys, its leaves are replaced one by one, and the call returns self. A leaf keeps its storage when the reshape can return a view of it. A lazy stack, or a tensordict that holds one, raises a NotImplementedError: call to_tensordict() first. See In-place operations. Defaults to False.

Returns:

A tensordict with the new batch size, or self when inplace=True or when the batch size does not change.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict({
...     'x': torch.arange(12).reshape(3, 4),
... }, batch_size=[3, 4])
>>> td = td.reshape(12)
>>> print(td['x'])
tensor([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11])

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

Tutorials that use reshape

Changing the shape of a TensorDict

Changing the shape of a TensorDict