For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Mojo function

softmax_inline

def softmax_inline[dtype: DType, simd_width: Int, rank: Int](input: TileTensor[dtype, Storage=input.Storage, address_space=input.address_space, linear_idx_type=input.linear_idx_type], output: TileTensor[dtype, Storage=output.Storage, address_space=output.address_space, linear_idx_type=output.linear_idx_type], axis: Int)

Computes softmax over the given axis of input and stores the result in output.

Wraps input with a load lambda and delegates to the main softmax_inline entry point.

Parameters:

• ​dtype (DType): The dtype of the input and output buffers.
• ​simd_width (Int): The simd_width to use in vectorization.
• ​rank (Int): The rank of the input and output tensors.

Args:

• ​input (TileTensor[dtype, Storage=input.Storage, address_space=input.address_space, linear_idx_type=input.linear_idx_type]): The input buffer used to compute the softmax.
• ​output (TileTensor[dtype, Storage=output.Storage, address_space=output.address_space, linear_idx_type=output.linear_idx_type]): The output buffer in which to store the softmax values.
• ​axis (Int): The axis along which to compute the softmax.

def softmax_inline[dtype: DType, simd_width: Int, rank: Int, input_fn: def[_simd_width: Int](Coord[*?]) capturing thin -> SIMD[dtype, _simd_width], target: StringSpan[ImmStaticOrigin] = StringSpan("cpu"), logsoftmax: Bool = False, has_prologue_fusion: Bool = True](shape: Coord, output: TileTensor[dtype, Storage=output.Storage, address_space=output.address_space, linear_idx_type=output.linear_idx_type], axis: Int, context: Optional[DeviceContext] = None)

Dispatches softmax (or logsoftmax) to the CPU or GPU target.

Selects the appropriate CPU or GPU implementation based on target and traces the operation. Exits early when the tensor is empty.

Parameters:

• ​dtype (DType): The dtype of the input and output buffers.
• ​simd_width (Int): The simd_width to use in vectorization.
• ​rank (Int): The rank of the input and output tensors.
• ​input_fn (def[_simd_width: Int](Coord[*?]) capturing thin -> SIMD[dtype, _simd_width]): The elementwise input lambda.
• ​target (StringSpan[ImmStaticOrigin]): The target device ("cpu" or "gpu").
• ​logsoftmax (Bool): Enable to apply elementwise log() to outputs after softmax.
• ​has_prologue_fusion (Bool): Whether the input lambda supports prologue fusion.

Args:

• ​shape (Coord): The shape of the output tensor.
• ​output (TileTensor[dtype, Storage=output.Storage, address_space=output.address_space, linear_idx_type=output.linear_idx_type]): The output buffer in which to store the softmax values.
• ​axis (Int): The axis along which to compute the softmax.
• ​context (Optional[DeviceContext]): Optional device context for GPU execution.