Function segmented_mm 

pub fn segmented_mm(
    a: impl AsRef<Array>,
    b: impl AsRef<Array>,
    segments: impl AsRef<Array>,
) -> Result<Array>

Perform a segmented matrix multiplication.

This computes multiple matrix multiplications where each segment of the reduction dimension is multiplied independently. This is useful for operations like mixture of experts or multi-head attention where different segments use different weights.

Params

• a: Input array with shape (M, K)
• b: Input array with shape (K, N)
• segments: Array of segment boundaries with shape (num_segments, 2). Each row contains [start, end) indices along the K dimension.

Returns

Array with shape (num_segments, M, N) where each segment contains the matrix multiplication for that segment of the K dimension.

Example

use mlx_rs::{Array, ops::segmented_mm};

let a = Array::ones::<f32>(&[10, 100]).unwrap();
let b = Array::ones::<f32>(&[100, 10]).unwrap();
let segments = Array::from_slice(&[0u32, 50, 50, 100], &[2, 2]);
let result = segmented_mm(&a, &b, &segments, None).unwrap();
// result has shape [2, 10, 10]