triton.language.softmax

triton.language.softmax(x, dim=None, *, keep_dims=None, ieee_rounding=False)

Computes the softmax of x along the given axis.

Parameters:

• x (Block) – the input values

• dim (int | None) – the axis along which to normalize. Defaults to 0 – note that this is not the last axis, unlike torch.softmax.

• keep_dims (bool | None) – deprecated and ignored. Softmax always preserves the input shape. Defaults to None; any other value emits a warning. Must be passed by keyword.

• ieee_rounding (bool) – whether the final division uses IEEE-compliant rounding. Must be passed by keyword.

This function can also be called as a member function on tensor, as x.softmax(...) instead of softmax(x, ...).