Function gather_mm 

pub fn gather_mm<'lhs, 'rhs>(
    a: impl AsRef<Array>,
    b: impl AsRef<Array>,
    lhs_indices: impl Into<Option<&'lhs Array>>,
    rhs_indices: impl Into<Option<&'rhs Array>>,
    sorted_indices: impl Into<Option<bool>>,
) -> Result<Array>

Matrix multiplication with gathered indices.

Perform matrix multiplication with index gathering along the batch dimensions. This is useful for operations where different batch elements should use different matrices from a pool.

Params

• a: Input array
• b: Input array
• lhs_indices: Optional indices to gather from a’s batch dimensions
• rhs_indices: Optional indices to gather from b’s batch dimensions
• sorted_indices: If true, indicates the indices are sorted which can enable optimizations (default: false)

Example

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

let a = Array::ones::<f32>(&[5, 32, 32]).unwrap();
let b = Array::ones::<f32>(&[3, 32, 32]).unwrap();
let lhs_indices = Array::from_slice(&[0u32, 2], &[2]);
let rhs_indices = Array::from_slice(&[2u32, 1], &[2]);
let result = gather_mm(&a, &b, &lhs_indices, &rhs_indices, None).unwrap();
// result has shape [2, 32, 32]