Function dispatch_softmax 

pub fn dispatch_softmax(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &DeviceRef,
    input: &MlxBuffer,
    output: &MlxBuffer,
    params_buf: &MlxBuffer,
    rows: u32,
    cols: u32,
) -> Result<()>

Dispatch a softmax operation on the GPU.

Arguments

• encoder - Command encoder to record the dispatch into.
• registry - Kernel registry (must have softmax sources registered).
• device - Metal device for pipeline compilation.
• input - Input buffer of shape [rows, cols] (f32, f16, or bf16).
• output - Output buffer (same dtype and shape as input).
• params_buf - Params buffer containing [cols, 0] as f32.
• rows - Number of rows.
• cols - Number of columns (softmax dimension).

Errors

Returns MlxError::InvalidArgument if:

• Input dtype is not f32, f16, or bf16.
• Input element count does not match rows * cols.