TensorDictBase.where#

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

Return a TensorDict of elements selected from either self or other, depending on condition.

Parameters:

condition (BoolTensor) – When True (nonzero), yields self, otherwise yields other.
other (TensorDictBase or Scalar) – value (if other is a scalar) or values selected at indices where condition is False. On a LazyStackedTensorDict, a tensor other that starts with the batch shape up to the stack dim is split along the stack dim, one piece per member, so it must also have the stack size at the stack dim: expand it to the batch shape first.

Keyword Arguments:

out (TensorDictBase, optional) – the output TensorDictBase instance.
pad (scalar, optional) – if provided, missing keys from the source or destination tensordict will be written as torch.where(mask, self, pad) or torch.where(mask, pad, other). Defaults to None, ie missing keys are not tolerated.
update_batch_size (bool, optional) – if True and out is provided, the batch size of the output will be updated to match the batch size of the condition. Defaults to False.

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

Docs