TensorDictBase.gather#
- TensorDictBase.gather(dim: int, index: Tensor, out: T | None = None, *, inplace: bool = False) Self#
Gathers values along a batch dimension, as
torch.gather()does.indexhas as many dimensions as the batch, and gives the batch size of the result; each leaf keeps its feature dimensions. As fortorch.gather(),indexcan be smaller than the batch along the other dimensions. The data is copied, and the names are kept.- Parameters:
dim (int) – the batch dimension to gather along.
index (Tensor) – a long tensor of indices along
dim, with as many dimensions as the batch.out (TensorDictBase, optional) – a tensordict to write the result into, with the batch size
index.shape. Exclusive withinplace. Defaults toNone.
- Keyword Arguments:
inplace (bool, optional) – if
True, the tensordict keeps its identity and keys, its leaves are replaced one by one, and the call returnsself. It needsindex.ndim == batch_dims. Exclusive without. A lazy stack, or a tensordict that holds one, raises aNotImplementedError: callto_tensordict()first. See In-place operations. Defaults toFalse.
Examples
>>> import torch >>> from tensordict import TensorDict >>> td = TensorDict( ... a=torch.arange(12).view(3, 4), ... b=TensorDict(c=torch.arange(12).view(3, 4, 1), batch_size=[3, 4, 1]), ... batch_size=[3, 4], ... ) >>> index = torch.tensor([[0, 3], [1, 2], [3, 3]]) >>> out = td.gather(dim=1, index=index) >>> out["a"] tensor([[ 0, 3], [ 5, 6], [11, 11]]) >>> out.batch_size, out["b", "c"].shape (torch.Size([3, 2]), torch.Size([3, 2, 1]))
The names of the batch dimensions are kept:
>>> td.names = ["x", "y"] >>> td.gather(dim=1, index=index).names ['x', 'y']
Every tensordict class and every tensorclass has this method. For the other methods, see Methods by task.
