TensorDictBase.where#

TensorDictBase.where(condition: Tensor, other: Tensor | TensorDictBase, *, out: TensorDictBase | None = None, pad: int | bool = None, update_batch_size: bool = False) → Self#

Returns a tensordict with the leaves where condition is True and other elsewhere.

torch.where() is applied to every leaf. condition is aligned with the batch dimensions of each leaf and expanded over its feature dimensions.

Parameters:

condition (Tensor) – a boolean mask over the leading batch dimensions.
other (TensorDictBase, Tensor or Number) – the values to take where condition is False. A tensordict is paired with the tensordict key by key. On a LazyStackedTensorDict, a tensor other that starts with the batch shape up to the stack dimension is split along that dimension, one piece per member, so it must have the stack size there: expand it to the batch shape first. The pieces line up with the leaves as in TensorDict only for the leaves with as many dims as other; for the others the lazy stack emits a FutureWarning: starting with TensorDict 0.17, the default will change to broadcasting like TensorDict.where, which aligns other with each leaf from the right.

Keyword Arguments:

out (TensorDictBase, optional) – the tensordict to write the result into, which is returned. Defaults to None, which returns a new tensordict.
pad (Number or bool, optional) – the value that stands in for an entry that only one of the tensordict and a tensordict other has. Defaults to None, which raises a KeyError for such an entry.
update_batch_size (bool, optional) – passed to update() when the result is written into out. If True, a lazy stack out gets as many members as the result. When other is not a tensordict, or self is a lazy stack, an out whose batch size is a prefix of the result’s ([] for [2]) keeps it, with a FutureWarning: starting with TensorDict 0.17, the default will change to giving out the batch size of the result, as update() does. Defaults to False.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(x=torch.ones(2, 3), batch_size=[2])
>>> td.where(torch.tensor([True, False]), 0.0)["x"]
tensor([[1., 1., 1.],
        [0., 0., 0.]])
>>> other = TensorDict(y=torch.zeros(2), batch_size=[2])
>>> sorted(td.where(torch.tensor([True, False]), other, pad=-1).keys())
['x', 'y']

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

TensorDictBase.where#

Docs

Tutorials

Resources