mean() and median() in PyTorch
Tags: pytorch, mean, median, function
Published: 2024-04-13 22:45:45

[mean()](https://pytorch.org/docs/stable/generated/torch.mean.html) can get the 0 or more D tensor of one or more mean(average) elements from the 0D or more D tensor of zero or more elements as shown below:

*Memos:
- `mean()` 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 `float` or `complex`).
- The 2nd argument with `torch` or the 1st argument with a tensor is `dim`(Optional-Type:`int`, `tuple` of `int` or `list` of `int`).
- The 3rd argument with `torch` or the 2nd argument with a tensor is `keepdim`(Optional-Default:`False`-Type:`bool`):
  *Memos:
  - It must be used with `dim`.
  - [My post](https://dev.to/hyperkai/set-keepdim-with-keepdim-argument-functions-pytorch-2fdj) explains `keepdim` argument.
- There is `dtype` argument with `torch` or a tesnor([dtype](https://pytorch.org/docs/stable/tensor_attributes.html#torch.dtype)) (Optional):
  *Memos:
  - If `dtype` is not given, `dtype` is inferred from `input` or `dtype` of [set_default_dtype()](https://pytorch.org/docs/stable/generated/torch.set_default_dtype.html) is used for floating-point numbers.
  - `dtype=` must be used.
  - [My post](https://dev.to/hyperkai/set-dtype-with-dtype-argument-functions-and-get-it-in-pytorch-13h2) explains `dtype` argument.
- There is `out` argument with `torch`(Optional-Type:`tensor`):
  *Memos:
  - It must be used with `dim`.
  - `out=` must be used.
  - [My post](https://dev.to/hyperkai/set-out-with-out-argument-functions-pytorch-3ee) explains `out` argument.
- Empty 1D or more D tensor gets a `nan`.

```python
import torch

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

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

my_tensor = torch.tensor([[5., 4., 7., 7.],
                          [6., 5., 3., 5.],
                          [3., 8., 9., 3.]])
torch.mean(input=my_tensor)
torch.mean(input=my_tensor, dim=(0, 1))
torch.mean(input=my_tensor, dim=(0, -1))
torch.mean(input=my_tensor, dim=(1, 0))
torch.mean(input=my_tensor, dim=(1, -2))
torch.mean(input=my_tensor, dim=(-1, 0))
torch.mean(input=my_tensor, dim=(-1, -2))
torch.mean(input=my_tensor, dim=(-2, 1))
torch.mean(input=my_tensor, dim=(-2, -1))
# tensor(5.4167)

torch.mean(input=my_tensor, dim=0)
torch.mean(input=my_tensor, dim=(0,))
torch.mean(input=my_tensor, dim=-2)
torch.mean(input=my_tensor, dim=(-2,))
# tensor([4.6667, 5.6667, 6.3333, 5.0000])

torch.mean(input=my_tensor, dim=1)
torch.mean(input=my_tensor, dim=(1,))
torch.mean(input=my_tensor, dim=-1)
torch.mean(input=my_tensor, dim=(-1,))
# tensor([5.7500, 4.7500, 5.7500])

my_tensor = torch.tensor([[5.+0.j, 4.+0.j, 7.+0.j, 7.+0.j],
                          [6.+0.j, 5.+0.j, 3.+0.j, 5.+0.j],
                          [3.+0.j, 8.+0.j, 9.+0.j, 3.+0.j]])
torch.mean(input=my_tensor)
# tensor(5.4167+0.j)

my_tensor = torch.tensor([])

torch.mean(input=my_tensor)
# tensor(nan)
```

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

*Memos:
- `median()` can be used with `torch` or a tensor.
- The 1st argument with `torch` or using a tensor is `input`(Required-Type:`tensor` of `int` or `float`).
- The 2nd argument with `torch` or the 1st argument with a tensor is `dim`(Optional-Type:`int`).
- The 3rd argument with `torch` or the 2nd argument with a tensor is `keepdim`(Optional-Default:`False`-Type:`bool`):
  *Memos:
  - It must be used with `dim`.
  - [My post](https://dev.to/hyperkai/set-keepdim-with-keepdim-argument-functions-pytorch-2fdj) explains `keepdim` argument.
- There is `out` argument with `torch`(Optional-Type:`tuple`(`tensor`, `tensor`) or `list`(`tensor`, `tensor`):
  *Memos:
  - It must be used with `dim`.
  - `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, 4, 7, 7])

torch.median(input=my_tensor)
my_tensor.median()
# tensor(5)

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

my_tensor = torch.tensor([[5, 4, 7, 7],
                          [6, 5, 3, 5],
                          [3, 8, 9, 3]])
torch.median(input=my_tensor)
# tensor(5)

torch.median(input=my_tensor, dim=0)
torch.median(input=my_tensor, dim=-2)
# torch.return_types.median(
# values=tensor([5, 5, 7, 5]),
# indices=tensor([0, 1, 0, 1]))

torch.median(input=my_tensor, dim=1)
torch.median(input=my_tensor, dim=-1)
# torch.return_types.median(
# values=tensor([5, 5, 3]),
# indices=tensor([0, 1, 3]))

my_tensor = torch.tensor([[5., 4., 7., 7.],
                          [6., 5., 3., 5.],
                          [3., 8., 9., 3.]])
torch.median(input=my_tensor)
# tensor(5.)

my_tensor = torch.tensor([])

torch.median(input=my_tensor)
# tensor(nan)
```