reference/generated/tensordict.utils.isin

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

isin

tensordict.utils.isin(input: TensorDictBase, reference: TensorDictBase, key: NestedKey, dim: int = 0) → Tensor

Returns which elements of input along dim are also in reference.

The entry key is cut into slices along dim. An element of input is in reference if its slice is equal to one of the slices of reference[key]. The two tensordicts must have the same number of batch dimensions, and the shapes of the two entries must match except along dim.

Parameters:

• input (TensorDictBase) – the tensordict to test.

• reference (TensorDictBase) – the tensordict to look in.

• key (NestedKey) – the key of the tensor to compare. It must hold a tensor in both tensordicts.

• dim (int, optional) – the batch dimension along which to compare. Defaults to 0.

Returns:

A boolean tensor of length input.batch_size[dim].

Raises:

• KeyError – if key is missing or is not a tensor in either tensordict.

• ValueError – if the numbers of batch dimensions differ, or if dim is not a batch dimension.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> from tensordict.utils import isin
>>> td = TensorDict(
...     x=torch.tensor([[1, 2], [3, 4], [1, 2], [5, 6]]),
...     y=torch.arange(4),
...     batch_size=[4],
... )
>>> reference = TensorDict(x=torch.tensor([[1, 2], [5, 6]]), batch_size=[2])
>>> isin(td, reference, key="x")
tensor([ True, False,  True,  True])