Function tropical_matmul_with_argmax 

pub fn tropical_matmul_with_argmax<T: TropicalWithArgmax<Index = u32> + KernelDispatch>(
    a: &[T::Scalar],
    m: usize,
    k: usize,
    b: &[T::Scalar],
    n: usize,
) -> GemmWithArgmax<T>

Tropical matrix multiplication with argmax tracking.

Returns both the result matrix and the argmax indices indicating which k produced each optimal C[i,j].

Example

use tropical_gemm::{tropical_matmul_with_argmax, TropicalMaxPlus};

let a = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0]; // 2x3
let b = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0]; // 3x2

let result = tropical_matmul_with_argmax::<TropicalMaxPlus<f64>>(&a, 2, 3, &b, 2);
assert_eq!(result.m, 2);
assert_eq!(result.n, 2);