Set keepdim with keepdim argument functions PyTorch
Tags: pytorch, keepdim, argument, function
Published: 2024-06-28 03:32:49

You can set `keepdim` with the functions which have `keepdim` argument as shown below:

*Memos:
- I selected some popular `keepdim` argument functions such as [sum()](https://pytorch.org/docs/stable/generated/torch.sum.html),
[prod()](https://pytorch.org/docs/stable/generated/torch.prod.html)
[mean()](https://pytorch.org/docs/stable/generated/torch.mean.html),
[median()](https://pytorch.org/docs/stable/generated/torch.median.html),
[min()](https://pytorch.org/docs/stable/generated/torch.min.html), [max()](https://pytorch.org/docs/stable/generated/torch.max.html), 
[argmin()](https://pytorch.org/docs/stable/generated/torch.argmin.html),
[argmax()](https://pytorch.org/docs/stable/generated/torch.argmax.html),
[all()](https://pytorch.org/docs/stable/generated/torch.all.html) and [any()](https://pytorch.org/docs/stable/generated/torch.any.html).
- `keepdim`(Optional-Default:`False`-Type:`bool`) can keep the dimension of `input` tensor.
- Sometimes, `keepdim` needs to be used with `dim`.

`sum()`. *[My post](https://dev.to/hyperkai/sum-prod-and-cartesianprod-in-pytorch-22li) explains `sum()`:

```python
import torch

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

torch.sum(input=my_tensor)
torch.sum(input=my_tensor, dim=0)
# tensor(10)

torch.sum(input=my_tensor, dim=0, keepdim=True)
# tensor([10])
```

`prod()`. *[My post](https://dev.to/hyperkai/sum-prod-and-cartesianprod-in-pytorch-22li) explains `prod()`:

```python
import torch

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

torch.prod(input=my_tensor)
torch.prod(input=my_tensor, dim=0)
# tensor(24)

torch.prod(input=my_tensor, dim=0, keepdim=True)
# tensor([24])
```

`mean()`. *[My post](https://dev.to/hyperkai/mean-and-median-in-pytorch-23h7) explains `mean()`:

```python
import torch

my_tensor = torch.tensor([5., 4., 7., 7.])

torch.mean(input=my_tensor)
torch.mean(input=my_tensor, dim=0)
# tensor(5.7500)

torch.mean(input=my_tensor, dim=0, keepdim=True)
tensor([5.7500])
```

`median()`. *[My post](https://dev.to/hyperkai/mean-and-median-in-pytorch-23h7) explains `median()`:

```python
import torch

my_tensor = torch.tensor([5, 4, 7, 7])

torch.median(input=my_tensor, dim=0)
# torch.return_types.median(
# values=tensor(5),
# indices=tensor(0))

torch.median(input=my_tensor, dim=0, keepdim=True)
# torch.return_types.median(
# values=tensor([5]),
# indices=tensor([0]))
```

`min()`. *[My post](https://dev.to/hyperkai/min-max-argmin-and-argmax-in-pytorch-2d6) explains `min()`:

```python
import torch

my_tensor = torch.tensor([5, 4, 7, 7])

torch.min(input=my_tensor, dim=0)
# torch.return_types.min(
# values=tensor(4),
# indices=tensor(1))

torch.min(input=my_tensor, dim=0, keepdim=True)
# torch.return_types.min(
# values=tensor([4]),
# indices=tensor([1]))
```

`max()`. *[My post](https://dev.to/hyperkai/min-max-argmin-and-argmax-in-pytorch-2d6) explains `max()`:

```python
import torch

my_tensor = torch.tensor([5, 4, 7, 7])

torch.max(input=my_tensor, dim=0)
# torch.return_types.max(
# values=tensor(7),
# indices=tensor(2))

torch.max(input=my_tensor, dim=0, keepdim=True)
# torch.return_types.max(
# values=tensor([7]),
# indices=tensor([2]))
```

`argmin()`. *[My post](https://dev.to/hyperkai/min-max-argmin-and-argmax-in-pytorch-2d6) explains `argmin()`:

```python
import torch

my_tensor = torch.tensor([5, 4, 7, 7])

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

torch.argmin(input=my_tensor, keepdim=True)
torch.argmin(input=my_tensor, dim=0, keepdim=True)
# tensor([1])
```

`argmax()`. *[My post](https://dev.to/hyperkai/min-max-argmin-and-argmax-in-pytorch-2d6) explains `argmax()`:

```python
import torch

my_tensor = torch.tensor([5, 4, 7, 7])

torch.argmax(input=my_tensor)
torch.argmax(input=my_tensor, dim=0)
# tensor(2)

torch.argmax(input=my_tensor, keepdim=True)
torch.argmax(input=my_tensor, dim=0, keepdim=True)
# tensor([2])
```

`all()`. *[My post](https://dev.to/hyperkai/min-max-argmin-and-argmax-in-pytorch-2d6) explains `all()`:

```python
import torch

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

torch.all(input=my_tensor)
torch.all(input=my_tensor, dim=0)
# tensor(False)

torch.all(input=my_tensor, keepdim=True)
torch.all(input=my_tensor, dim=0, keepdim=True)
# tensor([False])
```

`any()`. *[My post](https://dev.to/hyperkai/min-max-argmin-and-argmax-in-pytorch-2d6) explains `any()`:

```python
import torch

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

torch.any(input=my_tensor)
torch.any(input=my_tensor, dim=0)
# tensor(True)

torch.any(input=my_tensor, keepdim=True)
torch.any(input=my_tensor, dim=0, keepdim=True)
# tensor([True])
```