reference/generated/tensordict.TensorDictBase.mean

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.mean

TensorDictBase.mean(dim: int | Tuple[int] | Literal['feature'] = NO_DEFAULT, keepdim: bool = NO_DEFAULT, *, dtype: dtype | None = None, reduce: bool | None = None, key_transform: Callable[[NestedKey], NestedKey] | None = None) → Self | Tensor

Returns the mean of every leaf, over all its elements or along batch dimensions.

Parameters:

• dim (int, tuple of int or str, optional) – the batch dimensions to reduce. A negative dim counts from the end of the batch dimensions. "feature" reduces the feature dimensions of each leaf instead, and the result keeps the batch size. Defaults to None, which reduces every element of each leaf, and the result has an empty batch size.

• keepdim (bool, optional) – if True, the reduced batch dimensions are kept with size 1. Not allowed with dim="feature". Defaults to False.

Keyword Arguments:

• dtype (dtype, optional) – if given, each leaf is cast to this dtype before the reduction, which helps to prevent overflows. Defaults to None.

• reduce (bool, optional) – if True, reduces all the leaves together and returns a single tensor. Without dim, the result is a scalar. With an int dim, the leaves are concatenated along dim first, so they must have the same number of dimensions. With dim="feature", the result has the batch size as its shape. With feature dimensions, a negative dim counts from the end of each leaf. Starting with TensorDict 0.17, the default will change to counting it from the end of the batch dimensions, as with reduce=False. Defaults to False.

• key_transform (Callable, optional) – a function applied to the name of every entry of the result, at every level of nesting, such as lambda key: f"avg_{key}". Ignored when reduce=True. Defaults to None.

Returns:

A new tensordict, or a tensor if reduce=True.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> _ = torch.manual_seed(0)
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(c=torch.randn(3, 4, 5, 6), batch_size=(3, 4, 5)),
...     batch_size=(3, 4),
... )
>>> td.mean(dim=0).batch_size
torch.Size([4])
>>> td.mean(dim=0)["b", "c"].shape
torch.Size([4, 5, 6])
>>> td.mean()["a"].shape
torch.Size([])
>>> td.mean(dim="feature")["a"].shape
torch.Size([3, 4])
>>> td.mean(reduce=True)
tensor(0.0156)
>>> list(td.mean(key_transform=lambda key: f"avg_{key}").keys(True, True))
['avg_a', ('avg_b', 'avg_c')]

With reduce=True and an int dim, the leaves are concatenated along dim before the reduction:

>>> td = TensorDict(a=torch.ones(3, 4), b=torch.ones(3, 4), batch_size=(3, 4))
>>> td.mean(dim=0, reduce=True)
tensor([1., 1., 1., 1.])
>>> td.mean(dim="feature", reduce=True)
tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]])

Every tensordict class and every tensorclass has this method. For the other methods, see Methods by task.