Function transpose_2d 

pub fn transpose_2d(
    encoder: &mut CommandEncoder,
    registry: &mut KernelRegistry,
    device: &DeviceRef,
    input: &MlxBuffer,
    output: &MlxBuffer,
    rows: usize,
    cols: usize,
    dtype: DType,
) -> Result<()>

Encode a 2D matrix transpose: output[col, row] = input[row, col].

Buffer expectations

• input — [rows, cols] in the given dtype
• output — [cols, rows] in the given dtype (must be pre-allocated)

Errors

Returns MlxError::InvalidArgument if:

• rows or cols is zero
• dtype is not F32 or F16
• Buffers are too small