where¶

Documentation¶

treetensor.torch.where(condition, x, y)[source]¶

Return a tree of tensors of elements selected from either x or y, depending on condition.

Examples:

>>> import torch
>>> import treetensor.torch as ttorch
>>> ttorch.where(
...     torch.tensor([[True, False], [False, True]]),
...     torch.tensor([[2, 8], [16, 4]]),
...     torch.tensor([[3, 11], [5, 7]]),
... )
tensor([[ 2, 11],
        [ 5,  4]])

>>> tt1 = ttorch.randint(1, 99, {'a': (2, 3), 'b': {'x': (3, 2, 4)}})
>>> tt1
<Tensor 0x7f6760ad9908>
├── a --> tensor([[27, 90, 80],
│                 [12, 59,  5]])
└── b --> <Tensor 0x7f6760ad9860>
    └── x --> tensor([[[71, 52, 92, 79],
                       [48,  4, 13, 96]],

                      [[72, 89, 44, 62],
                       [32,  4, 29, 76]],

                      [[ 6,  3, 93, 89],
                       [44, 89, 85, 90]]])
>>> ttorch.where(tt1 % 2 == 1, tt1, 0)
<Tensor 0x7f6760ad9d30>
├── a --> tensor([[27,  0,  0],
│                 [ 0, 59,  5]])
└── b --> <Tensor 0x7f6760ad9f98>
    └── x --> tensor([[[71,  0,  0, 79],
                       [ 0,  0, 13,  0]],

                      [[ 0, 89,  0,  0],
                       [ 0,  0, 29,  0]],

                      [[ 0,  3, 93, 89],
                       [ 0, 89, 85,  0]]])

Torch Version Related

This documentation is based on torch.where in torch v2.4.1+cu121. Its arguments’ arrangements depend on the version of pytorch you installed.

If some arguments listed here are not working properly, please check your pytorch’s version with the following command and find its documentation.

python -c 'import torch;print(torch.__version__)'

The arguments and keyword arguments supported in torch v2.4.1+cu121 is listed below.

Description From Torch v2.4.1+cu121¶

torch.where(condition, x, y) → Tensor ¶

Return a tensor of elements selected from either x or y, depending on condition.

The operation is defined as:

\[\begin{split}\text{out}_i = \begin{cases} \text{x}_i & \text{if } \text{condition}_i \\ \text{y}_i & \text{otherwise} \\ \end{cases}\end{split}\]

Note

The tensors condition, x, y must be broadcastable.

Note

Currently valid scalar and tensor combination are 1. Scalar of floating dtype and torch.double 2. Scalar of integral dtype and torch.long 3. Scalar of complex dtype and torch.complex128

Arguments:

condition (BoolTensor): When True (nonzero), yield x, otherwise yield y x (Tensor or Scalar): value (if x is a scalar) or values selected at indices

where condition is True

y (Tensor or Scalar): value (if y is a scalar) or values selected at indices: where condition is False

Returns:

Tensor: A tensor of shape equal to the broadcasted shape of condition, x, y

Example:

>>> x = torch.randn(3, 2)
>>> y = torch.ones(3, 2)
>>> x
tensor([[-0.4620,  0.3139],
        [ 0.3898, -0.7197],
        [ 0.0478, -0.1657]])
>>> torch.where(x > 0, x, y)
tensor([[ 1.0000,  0.3139],
        [ 0.3898,  1.0000],
        [ 0.0478,  1.0000]])
>>> x = torch.randn(2, 2, dtype=torch.double)
>>> x
tensor([[ 1.0779,  0.0383],
        [-0.8785, -1.1089]], dtype=torch.float64)
>>> torch.where(x > 0, x, 0.)
tensor([[1.0779, 0.0383],
        [0.0000, 0.0000]], dtype=torch.float64)

torch.where(condition) → tuple of LongTensor

torch.where(condition) is identical to torch.nonzero(condition, as_tuple=True).

Note