reference/generated/tensordict.TensorDictBase.nanmean

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.nanmean

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

Returns the mean of all non-NaN elements in the input tensordict.

Parameters:

• dim (int, tuple of int, optional) – if None, returns a dimensionless tensordict containing the mean value of all leaves (if this can be computed). If integer or tuple of integers, mean is called upon the dimension specified if and only if this dimension is compatible with the tensordict shape. Only the “feature” string is currently permitted. Using dim=”feature” will achieve the reduction over all feature dimensions. If reduce=True, a tensor of the shape of the TensorDict’s batch-size will be returned. Otherwise, a new tensordict with the same structure as self with reduced feature dimensions will be returned.

• keepdim (bool) – whether the output tensor has dim retained or not.

Keyword Arguments:

• dtype (dtype, optional) – the desired data type of returned tensor. If specified, the input tensor is cast to dtype before the operation is performed. This is useful for preventing data type overflows. Default: None.

• reduce (bool, optional) – if True, the reduction will occur across all TensorDict values and a single reduced tensor will be returned. Defaults to False. With feature dims, 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 dims, as with reduce=False.

Examples

>>> from tensordict import TensorDict
>>> import torch
>>> _ = torch.manual_seed(0)
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.nanmean(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.nanmean()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> td.nanmean(reduce=True)
tensor(0.0456)
>>> td.nanmean(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.nanmean(reduce=True, dim="feature")
tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]])
>>> td.nanmean(reduce=True, dim=0)
tensor([[1., 1., 1., 1., 1.],
        [1., 1., 1., 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.