reference/generated/tensordict.TensorDictBase.softmax

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.softmax

TensorDictBase.softmax(dim: int, dtype: dtype | None = None)

Returns a new tensordict with torch.softmax() applied to every leaf along a batch dimension.

A LazyStackedTensorDict, also a nested one, gives a lazy stack with the values of its dense copy. The exception is a softmax along the stack dim of members whose entries have different shapes, which have no dense copy: each leaf of each member is normalised along its own dim dim.

Parameters:

• dim (int) – the batch dimension along which softmax is computed. A negative dim counts from the end of the batch dimensions.

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

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(x=torch.zeros(2, 3), batch_size=[2])
>>> td.softmax(0)["x"]
tensor([[0.5000, 0.5000, 0.5000],
        [0.5000, 0.5000, 0.5000]])

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