argwhere() and nonzero() in PyTorch
Tags: pytorch, argwhere, nonzero, indice
Published: 2024-05-16 18:50:31

*[My post](https://dev.to/hyperkai/where-and-countnonzero-in-pytorch-195d) explains [where()](https://pytorch.org/docs/stable/generated/torch.where.html) and [count_nonzero()](https://pytorch.org/docs/stable/generated/torch.count_nonzero.html).

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

*Memos:
- `argwhere()` can be used with [torch](https://pytorch.org/docs/stable/torch.html) or a tensor.
- The 1st argument with `torch` or using a tensor is `input`(Required-Type:`tensor` of `int`, `float`, `complex` or `bool`).

```python
import torch

my_tensor = torch.tensor(5)

torch.argwhere(input=my_tensor)
my_tensor.argwhere()
# tensor([], size=(1, 0), dtype=torch.int64)

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

torch.argwhere(input=my_tensor)
# tensor([[0], [2], [4], [5]])

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

torch.argwhere(input=my_tensor)
# tensor([[0], [2], [4], [5]])

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

torch.argwhere(input=my_tensor)
# tensor([[0], [2], [4], [5]])

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

torch.argwhere(input=my_tensor)
# tensor([[0], [2], [4]])

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

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

[nonzero()](https://pytorch.org/docs/stable/generated/torch.nonzero.html) can get the 2D tensor of the zero or more indices of non-zero elements or the one or more 1D tensors of the zero or more indices of non-zero elements from the 0D or more D tensor of zero or more elements as shown below:

*Memos:
- `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`).
- There is `as_tuple` argument with `torch` or a tensor(Optional-Default:`False`-Type:`bool`):
  *Memos: 
  - If `as_tuple` is `True`, the tuple of zero or more 1D tensors is returned.
- 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.

```python
import torch

my_tensor = torch.tensor(5)

torch.nonzero(input=my_tensor)
my_tensor.nonzero()
# tensor([], size=(1, 0), dtype=torch.int64)

torch.nonzero(input=my_tensor, as_tuple=True)
# (tensor([0]),)

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

torch.nonzero(input=my_tensor)
my_tensor.nonzero()
# tensor([[0], [2], [4], [5]])

torch.nonzero(input=my_tensor, as_tuple=True)
# (tensor([0, 2, 4, 5]),)

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

torch.nonzero(input=my_tensor)
# tensor([[0], [2], [4], [5]])

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

torch.nonzero(input=my_tensor)
# tensor([[0], [2], [4], [5]])

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

torch.nonzero(input=my_tensor)
# tensor([[0], [2], [4]])

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

torch.nonzero(input=my_tensor, as_tuple=True)
# (tensor([0, 0, 1, 1]),
#  tensor([0, 2, 1, 2]))

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

torch.nonzero(input=my_tensor, as_tuple=True)
# (tensor([0, 0, 0, 0, 1, 1, 1]),
#  tensor([0, 0, 1, 1, 0, 1, 1]),
#  tensor([0, 2, 1, 2, 1, 1, 2]))
```