torchoutil.nn.functional.activation module

log_softmax_multidim(
x: Tensor,
*,
dims: Iterable[int] | None = (-1,),
) Tensor[source]

A multi-dimensional version of torch.log_softmax along multiple dimensions at the same time.

softmax_multidim(
x: Tensor,
*,
dims: Iterable[int] | None = (-1,),
) Tensor[source]

A multi-dimensional version of torch.softmax along multiple dimensions at the same time.