where() and count_nonzero() in PyTorch
Tags: pytorch, where, countnonzero, function
Published: 2024-05-16 18:54:32

*[My post](https://dev.to/hyperkai/argwhere-and-nonzero-in-pytorch-1oeb) explains [argwhere()](https://pytorch.org/docs/stable/generated/torch.argwhere.html) and [nonzero()](https://pytorch.org/docs/stable/generated/torch.nonzero.html).

[where()](https://pytorch.org/docs/stable/generated/torch.where.html) can get the 0D or more D tensor of zero or more elements from two of the 0D or more D tensors of the zero or more elements selected either from `input` or `other` tensor, depending on `condition` as shown below:

*Memos:
- `where()` can be used with [torch](https://pytorch.org/docs/stable/torch.html) or a tensor.
- The 1st argument with `torch` or a tensor is `condition`(Required-Type:`tensor` of `bool`).
- The 2nd argument with `torch` or using a tensor is `input`(Required-Type:`tensor` or `scalar` of `int`, `float`, `complex` or `bool`):
  *Memos:
  - `torch` must use `input` with a scalar without `condition=`, `input=` and `other=`.
  - A tensor cannot use `input` with a scalar.
- The 3rd argument with `torch` or the 2nd argument with a tensor is `other`(Required-Type:`tensor` or `scalar` of `int`, `float`, `complex` or `bool`).
- There is `out` argument with `torch`(Optional-Type:`tensor`):
  *Memos:
  - `out=` must be used.
  - [My post](https://dev.to/hyperkai/set-out-with-out-argument-functions-pytorch-3ee) explains `out` argument.
- If `condition` is `True`, the element of `input` is selected otherwise the element of `other` is selected.

```python
import torch

tensor1 = torch.tensor([[5, 0, 4],
                        [0, 3, 1]])
tensor2 = torch.tensor([60, 70, 80])

torch.where(condition=tensor1 > 2, input=tensor1, other=tensor2)
tensor1.where(condition=tensor1 > 2, other=tensor2)
# tensor([[5, 70, 4],
#         [60, 3, 80]])

torch.where(condition=tensor1 > 2, input=tensor2, other=tensor1)
# tensor([[60, 0, 80],
#         [0, 70, 1]])

torch.where(tensor1 > 2, 10, tensor2)
# tensor([[10, 70, 10],
#         [60, 10, 80]])

torch.where(condition=tensor1 > 2, input=tensor1, other=10)
# tensor([[5, 10, 4],
#         [10, 3, 10]])

torch.where(tensor1 > 2, 10, 20)
# tensor([[10, 20, 10],
#         [20, 10, 20]])

tensor1 = torch.tensor([[5., 0., 4.],
                        [0., 3., 1.]])
tensor2 = torch.tensor([60., 70., 80.])
tensor3 = torch.tensor(True)
torch.where(condition=tensor3, input=tensor1, other=tensor2)
# tensor([[5., 0., 4.],
#         [0., 3., 1.]])

torch.where(tensor3, 5., other=tensor2)
# tensor([5., 5., 5.])

torch.where(condition=tensor3, input=tensor1, other=60.)
# tensor([[5., 0., 4.],
#         [0., 3., 1.]])

torch.where(tensor3, 5., other=60.)
# tensor(5.)

tensor1 = torch.tensor([[5.+0.j, 0.+0.j, 4.+0.j],
                        [0.+0.j, 3.+0.j, 1.+0.j]])
tensor2 = torch.tensor([60.+0.j, 70.+0.j, 80.+0.j])
tensor3 = torch.tensor(False)
torch.where(condition=tensor3, input=tensor1, other=tensor2)
# tensor([[60.+0.j, 70.+0.j, 80.+0.j],
#         [60.+0.j, 70.+0.j, 80.+0.j]])

torch.where(tensor3, 5.+0.j, other=tensor2)
# tensor([60.+0.j, 70.+0.j, 80.+0.j])

torch.where(condition=tensor3, input=tensor1, other=60.+0.j)
# tensor([[60.+0.j, 60.+0.j, 60.+0.j],
#         [60.+0.j, 60.+0.j, 60.+0.j]])

torch.where(tensor3, 5.+0.j, other=60.+0.j)
# tensor(60.+0.j)

tensor1 = torch.tensor([[True, False, True],
                        [False, True, False]])
tensor2 = torch.tensor([False, True, False])
tensor3 = torch.tensor(True)
torch.where(condition=tensor3, input=tensor1, other=tensor2)
# tensor([[True, False, True],
#         [False, True, False]])

torch.where(tensor3, True, other=tensor2)
# tensor([True, True, True])

torch.where(condition=tensor3, input=tensor1, other=False)
# tensor([[True, False, True],
#         [False, True, False]])

torch.where(tensor3, True, other=False)
# tensor(True)

tensor1 = torch.tensor([[[5, 0, 4], [0, 3, 1]],
                          [[0, 7, 0], [0, 6, 8]]])
tensor2 = torch.tensor([60, 70, 80])

torch.where(condition=tensor1 > 2, input=tensor1, other=tensor2)
# tensor([[[5, 70, 4],
#          [60, 3, 80]],
#         [[60, 7, 80],
#          [60, 6, 8]]])

torch.where(condition=tensor1 > 2, input=tensor2, other=tensor1)
# tensor([[[60, 0, 80],
#          [0, 70, 1]],
#         [[0, 70, 0],
#          [0, 70, 80]]])

torch.where(tensor1 > 2, 10, tensor2)
# tensor([[[10, 70, 10],
#          [60, 10, 80]],
#         [[60, 10, 80],
#          [60, 10, 10]]])

torch.where(condition=tensor1 > 2, input=tensor1, other=10)
# tensor([[[5, 10, 4],
#          [10, 3, 10]],
#         [[10, 7, 10],
#          [10, 6, 8]]])

torch.where(tensor1 > 2, 10, 20)
# tensor([[[10, 20, 10],
#          [20, 10, 20]],
#         [[20, 10, 20],
#          [20, 10, 10]]])
```

[count_nonzero()](https://pytorch.org/docs/stable/generated/torch.count_nonzero.html) can get the 0D or more D tensor of the count number of zero or more non-zero elements from the 0D or more D tensor of zero or more elements as shown below:

*Memos:
- `count_nonzero()` can be used with `torch` or a tensor.
- The 1st argument with `torch` or using a tensor is `input`(Required-Type:`tensor` of `int`, `float`, `complex` or `bool`).
- The 2nd argument with `torch` or the 1st argument with a tensor is `dim`(Optional-Type:`int`, `tuple` of `int` or `list` of `int`).

```python
import torch

my_tensor = torch.tensor(5)

torch.count_nonzero(input=my_tensor)
my_tensor.count_nonzero()
torch.count_nonzero(input=my_tensor, dim=0)
torch.count_nonzero(input=my_tensor, dim=-1)
torch.count_nonzero(input=my_tensor, dim=(0,))
torch.count_nonzero(input=my_tensor, dim=(-1,))
# tensor(1)

my_tensor = torch.tensor([5, 0, 4, 0, 3, 1])

torch.count_nonzero(input=my_tensor)
torch.count_nonzero(input=my_tensor, dim=0)
torch.count_nonzero(input=my_tensor, dim=-1)
torch.count_nonzero(input=my_tensor, dim=(0,))
torch.count_nonzero(input=my_tensor, dim=(-1,))
# tensor(4)

my_tensor = torch.tensor([5., 0., 4., 0., 3., 1.])

torch.count_nonzero(input=my_tensor)
torch.count_nonzero(input=my_tensor, dim=0)
torch.count_nonzero(input=my_tensor, dim=-1)
torch.count_nonzero(input=my_tensor, dim=(0,))
torch.count_nonzero(input=my_tensor, dim=(-1,))
# tensor(4)

my_tensor = torch.tensor([5.+0.j, 0.+0.j, 4.+0.j, 0.+0.j, 3.+0.j, 1.+0.j])

torch.count_nonzero(input=my_tensor)
# tensor(4)

my_tensor = torch.tensor([True, False, True, False, True, False])

torch.count_nonzero(input=my_tensor)
# tensor(3)

my_tensor = torch.tensor([[5, 0, 4],
                          [0, 3, 1]])
torch.count_nonzero(input=my_tensor)
torch.count_nonzero(input=my_tensor, dim=(0, 1))
torch.count_nonzero(input=my_tensor, dim=(0, -1))
torch.count_nonzero(input=my_tensor, dim=(1, 0))
torch.count_nonzero(input=my_tensor, dim=(1, -2))
torch.count_nonzero(input=my_tensor, dim=(-1, 0))
torch.count_nonzero(input=my_tensor, dim=(-1, -2))
torch.count_nonzero(input=my_tensor, dim=(-2, 1))
torch.count_nonzero(input=my_tensor, dim=(-2, -1))
# tensor(4)

torch.count_nonzero(input=my_tensor, dim=0)
torch.count_nonzero(input=my_tensor, dim=-2)
torch.count_nonzero(input=my_tensor, dim=(0,))
torch.count_nonzero(input=my_tensor, dim=(-2,))
# tensor([1, 1, 2])

torch.count_nonzero(input=my_tensor, dim=1)
torch.count_nonzero(input=my_tensor, dim=-1)
torch.count_nonzero(input=my_tensor, dim=(1,))
torch.count_nonzero(input=my_tensor, dim=(-1,))
# tensor([2, 2])

my_tensor = torch.tensor([[[5, 0, 4], [0, 3, 1]],
                          [[0, 7, 0], [0, 6, 8]]])
torch.count_nonzero(input=my_tensor)
# tensor(7)

torch.count_nonzero(input=my_tensor, dim=0)
torch.count_nonzero(input=my_tensor, dim=-3)
torch.count_nonzero(input=my_tensor, dim=(0,))
torch.count_nonzero(input=my_tensor, dim=(-3,))
# tensor([[1, 1, 1], [0, 2, 2]])

torch.count_nonzero(input=my_tensor, dim=1)
torch.count_nonzero(input=my_tensor, dim=-2)
torch.count_nonzero(input=my_tensor, dim=(1,))
torch.count_nonzero(input=my_tensor, dim=(-2,))
# tensor([[1, 1, 2], [0, 2, 1]])

torch.count_nonzero(input=my_tensor, dim=2)
torch.count_nonzero(input=my_tensor, dim=-1)
torch.count_nonzero(input=my_tensor, dim=(2,))
torch.count_nonzero(input=my_tensor, dim=(-1,))
# tensor([[2, 2], [1, 2]])

torch.count_nonzero(input=my_tensor, dim=(0, 1))
torch.count_nonzero(input=my_tensor, dim=(0, -2))
torch.count_nonzero(input=my_tensor, dim=(1, 0))
torch.count_nonzero(input=my_tensor, dim=(1, -3))
torch.count_nonzero(input=my_tensor, dim=(-2, 0))
torch.count_nonzero(input=my_tensor, dim=(-2, -3))
torch.count_nonzero(input=my_tensor, dim=(-3, 1))
torch.count_nonzero(input=my_tensor, dim=(-3, -2))
# tensor([1, 3, 3])

torch.count_nonzero(input=my_tensor, dim=(0, 2))
torch.count_nonzero(input=my_tensor, dim=(0, -1))
torch.count_nonzero(input=my_tensor, dim=(2, 0))
torch.count_nonzero(input=my_tensor, dim=(2, -3))
torch.count_nonzero(input=my_tensor, dim=(-1, 0))
torch.count_nonzero(input=my_tensor, dim=(-1, -3))
torch.count_nonzero(input=my_tensor, dim=(-3, 2))
torch.count_nonzero(input=my_tensor, dim=(-3, -1))
# tensor([3, 4])

torch.count_nonzero(input=my_tensor, dim=(1, 2))
torch.count_nonzero(input=my_tensor, dim=(1, -1))
torch.count_nonzero(input=my_tensor, dim=(2, 1))
torch.count_nonzero(input=my_tensor, dim=(2, -2))
torch.count_nonzero(input=my_tensor, dim=(-1, 1))
torch.count_nonzero(input=my_tensor, dim=(-1, -2))
torch.count_nonzero(input=my_tensor, dim=(-2, 2))
torch.count_nonzero(input=my_tensor, dim=(-2, -1))
# tensor([4, 3])
```