Mojo function

softmax_3_pass

softmax_3_pass[simd_width: Int, buffer_size: Dim, dtype: DType, origins: OriginSet, input_fn_1d: fn[_simd_width: Int](Int) capturing -> SIMD[dtype, _simd_width], logsoftmax: Bool = False](output: NDBuffer[dtype, 1, origin, DimList.__init__[Dim](buffer_size)])

Performs an unbatched softmax on an input tensor using the three-pass algorithm.

The unbatched three-pass softmax is defined as:

procedure SoftmaxUnbatched(InputInput)
  maxVal = -∞
  denom = 0
  STEP 1: find the max value in each batch
  for i = 0 to N do
    maxVal = max(maxVal, Input[b, i])
  end for
  STEP 2: compute the exponential for each batch
  for i = 0 to N do
    Output[b, i] = exp(Input[b, i] - maxVal)
    denom += Output[b, i]
  end for
  STEP 3: normalize each batch
  for i = 0 to N do
    Output[b, i] /= denom
  end for

Parameters:

simd_width (Int): The simd_width to use in vectorization.
buffer_size (Dim): The size of the input and output buffers.
dtype (DType): The dtype of the input and output buffers.
origins (OriginSet): The OriginSet of captured arguments by the input_fn_1d.
input_fn_1d (fn[_simd_width: Int](Int) capturing -> SIMD[dtype, _simd_width]): The elementwise input lambda.
logsoftmax (Bool): Enable to apply elementwise log() to outputs after softmax.

Args:

output (NDBuffer): The output buffer in which to store the softmax values.