asc2.softmax

asc2.softmax(input: LocalTensor) → LocalTensor

Computes the row-wise softmax of input.

For 2D tensors, softmax is applied independently along the last dimension for each row. For 1D tensors, softmax is applied over all elements.

The supported data types for the input are: float16, float32.

Parameters:

input – The input tensor (1D or 2D)

Returns:

The result tensor with the same shape as input

Return type:

LocalTensor

Raises:

RuntimeError – If the input dtype is not supported or input has more than 2 dimensions

Examples

Compute row-wise softmax for a 2D tensor:

input = asc2.copy_in(x_gm, [0, 0], [64, 1024])
result = asc2.softmax(input)  # softmax applied independently to each of 64 rows

Compute softmax for a 1D tensor:

input = asc2.copy_in(x_gm, [0], [1024])
result = asc2.softmax(input)  # softmax applied over all 1024 elements

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