tropical_gemm/mat/
mod.rs

1//! Matrix types for tropical algebra.
2//!
3//! This module provides faer-inspired matrix types:
4//! - [`Mat<S>`]: Owned matrix storing semiring values
5//! - [`MatRef<'a, S>`]: Immutable view over scalar data
6//! - [`MatMut<'a, S>`]: Mutable view over semiring data
7//!
8//! # Example
9//!
10//! ```
11//! use tropical_gemm::{Mat, MatRef, MaxPlus};
12//!
13//! // Create a view from raw data
14//! let data = [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0];
15//! let a = MatRef::<MaxPlus<f32>>::from_slice(&data, 2, 3);
16//! let b = MatRef::<MaxPlus<f32>>::from_slice(&data, 3, 2);
17//!
18//! // Matrix multiplication using method
19//! let c = a.matmul(&b);
20//!
21//! // Or using operator syntax
22//! let c = &a * &b;
23//!
24//! // Factory methods
25//! let zeros = Mat::<MaxPlus<f32>>::zeros(3, 3);
26//! let identity = Mat::<MaxPlus<f32>>::identity(3);
27//! ```
28
29mod mut_;
30mod ops;
31mod owned;
32mod ref_;
33
34pub use mut_::MatMut;
35pub use owned::Mat;
36pub use ref_::MatRef;
37
38/// Result of matrix multiplication with argmax tracking.
39pub struct MatWithArgmax<S: crate::TropicalWithArgmax> {
40    /// The result matrix values.
41    pub values: Mat<S>,
42    /// The argmax indices (which k produced each C[i,j]).
43    pub argmax: Vec<u32>,
44}
45
46impl<S: crate::TropicalWithArgmax<Index = u32>> MatWithArgmax<S> {
47    /// Get the value at position (i, j).
48    pub fn get(&self, i: usize, j: usize) -> S {
49        self.values[(i, j)]
50    }
51
52    /// Get the scalar value at position (i, j).
53    ///
54    /// This is a convenience method that extracts the underlying scalar
55    /// without requiring a trait import.
56    #[inline]
57    pub fn get_value(&self, i: usize, j: usize) -> S::Scalar {
58        self.values[(i, j)].value()
59    }
60
61    /// Get the argmax index at position (i, j).
62    pub fn get_argmax(&self, i: usize, j: usize) -> u32 {
63        // Column-major indexing
64        self.argmax[j * self.values.nrows() + i]
65    }
66
67    /// Number of rows.
68    pub fn nrows(&self) -> usize {
69        self.values.nrows()
70    }
71
72    /// Number of columns.
73    pub fn ncols(&self) -> usize {
74        self.values.ncols()
75    }
76
77    /// Get the argmax indices as a slice.
78    ///
79    /// This is useful for backward pass computation.
80    #[inline]
81    pub fn argmax_slice(&self) -> &[u32] {
82        &self.argmax
83    }
84
85    /// Compute gradient with respect to matrix A.
86    ///
87    /// Given the upstream gradient dL/dC, computes dL/dA using the argmax
88    /// indices from the forward pass.
89    ///
90    /// For C = A ⊗ B where C[i,j] = ⊕_k (A[i,k] ⊗ B[k,j]):
91    /// dL/dA[i,k] = Σ_j { dL/dC[i,j] if argmax[i,j] == k }
92    ///
93    /// # Arguments
94    ///
95    /// * `grad_c` - Gradient of the loss with respect to C, dimensions m×n
96    /// * `k` - Number of columns in A (the inner dimension)
97    ///
98    /// # Returns
99    ///
100    /// Gradient of the loss with respect to A, dimensions m×k
101    ///
102    /// # Example
103    ///
104    /// ```
105    /// use tropical_gemm::{Mat, MaxPlus, TropicalMaxPlus};
106    ///
107    /// let a = Mat::<MaxPlus<f64>>::from_col_major(&[1.0, 4.0, 2.0, 5.0, 3.0, 6.0], 2, 3);
108    /// let b = Mat::<MaxPlus<f64>>::from_col_major(&[1.0, 3.0, 5.0, 2.0, 4.0, 6.0], 3, 2);
109    ///
110    /// // Forward pass with argmax
111    /// let result = a.matmul_argmax(&b);
112    ///
113    /// // Backward pass: grad_c is upstream gradient (e.g., all ones)
114    /// let grad_c = Mat::<MaxPlus<f64>>::from_fn(2, 2, |_, _| TropicalMaxPlus(1.0));
115    /// let grad_a = result.backward_a(&grad_c, 3); // k=3 (columns in A)
116    ///
117    /// assert_eq!(grad_a.nrows(), 2);
118    /// assert_eq!(grad_a.ncols(), 3);
119    /// ```
120    /// Only MaxPlus/MinPlus support routing without the original inputs.
121    /// MaxMul requires `backward_a_maxmul` instead:
122    ///
123    /// ```compile_fail
124    /// use tropical_gemm::{Mat, MaxMul};
125    /// let a = Mat::<MaxMul<f64>>::from_col_major(&[2.0], 1, 1);
126    /// let result = a.matmul_argmax(&a);
127    /// let _ = result.backward_a(&a, 1);
128    /// ```
129    pub fn backward_a<G>(&self, grad_c: &Mat<G>, k: usize) -> Mat<G>
130    where
131        S: crate::types::AdditiveTropical,
132        G: crate::TropicalSemiring,
133        G::Scalar: Copy + Default + std::ops::AddAssign,
134    {
135        let m = self.nrows();
136        let n = self.ncols();
137        assert_eq!(grad_c.nrows(), m, "grad_c rows mismatch");
138        assert_eq!(grad_c.ncols(), n, "grad_c cols mismatch");
139
140        // Output is m×k in column-major
141        let mut grad_a_data = vec![G::Scalar::default(); m * k];
142
143        for j in 0..n {
144            for i in 0..m {
145                // Column-major indexing for argmax
146                let idx = self.argmax[j * m + i] as usize;
147                if idx < k {
148                    // Column-major indexing for grad_a: element (i, idx) at idx * m + i
149                    grad_a_data[idx * m + i] += grad_c[(i, j)].value();
150                }
151            }
152        }
153
154        Mat::from_col_major(&grad_a_data, m, k)
155    }
156
157    /// Compute gradient with respect to matrix B.
158    ///
159    /// Given the upstream gradient dL/dC, computes dL/dB using the argmax
160    /// indices from the forward pass.
161    ///
162    /// For C = A ⊗ B where C[i,j] = ⊕_k (A[i,k] ⊗ B[k,j]):
163    /// dL/dB[k,j] = Σ_i { dL/dC[i,j] if argmax[i,j] == k }
164    ///
165    /// # Arguments
166    ///
167    /// * `grad_c` - Gradient of the loss with respect to C, dimensions m×n
168    /// * `k` - Number of rows in B (the inner dimension)
169    ///
170    /// # Returns
171    ///
172    /// Gradient of the loss with respect to B, dimensions k×n
173    ///
174    /// # Example
175    ///
176    /// ```
177    /// use tropical_gemm::{Mat, MaxPlus, TropicalMaxPlus};
178    ///
179    /// let a = Mat::<MaxPlus<f64>>::from_col_major(&[1.0, 4.0, 2.0, 5.0, 3.0, 6.0], 2, 3);
180    /// let b = Mat::<MaxPlus<f64>>::from_col_major(&[1.0, 3.0, 5.0, 2.0, 4.0, 6.0], 3, 2);
181    ///
182    /// // Forward pass with argmax
183    /// let result = a.matmul_argmax(&b);
184    ///
185    /// // Backward pass: grad_c is upstream gradient
186    /// let grad_c = Mat::<MaxPlus<f64>>::from_fn(2, 2, |_, _| TropicalMaxPlus(1.0));
187    /// let grad_b = result.backward_b(&grad_c, 3); // k=3 (rows in B)
188    ///
189    /// assert_eq!(grad_b.nrows(), 3);
190    /// assert_eq!(grad_b.ncols(), 2);
191    /// ```
192    pub fn backward_b<G>(&self, grad_c: &Mat<G>, k: usize) -> Mat<G>
193    where
194        S: crate::types::AdditiveTropical,
195        G: crate::TropicalSemiring,
196        G::Scalar: Copy + Default + std::ops::AddAssign,
197    {
198        let m = self.nrows();
199        let n = self.ncols();
200        assert_eq!(grad_c.nrows(), m, "grad_c rows mismatch");
201        assert_eq!(grad_c.ncols(), n, "grad_c cols mismatch");
202
203        // Output is k×n in column-major
204        let mut grad_b_data = vec![G::Scalar::default(); k * n];
205
206        for j in 0..n {
207            for i in 0..m {
208                // Column-major indexing for argmax
209                let idx = self.argmax[j * m + i] as usize;
210                if idx < k {
211                    // Column-major indexing for grad_b: element (idx, j) at j * k + idx
212                    grad_b_data[j * k + idx] += grad_c[(i, j)].value();
213                }
214            }
215        }
216
217        Mat::from_col_major(&grad_b_data, k, n)
218    }
219}
220
221impl<T: crate::types::TropicalScalar> MatWithArgmax<crate::TropicalMaxMul<T>> {
222    /// MaxMul gradient with respect to A, using the winning values from B.
223    pub fn backward_a_maxmul<G>(&self, grad_c: &Mat<G>, b: &Mat<crate::TropicalMaxMul<T>>) -> Mat<G>
224    where
225        G: crate::TropicalSemiring<Scalar = T>,
226        T: Default + std::ops::AddAssign,
227    {
228        let (m, n, k) = (self.nrows(), self.ncols(), b.nrows());
229        assert_eq!(
230            (grad_c.nrows(), grad_c.ncols()),
231            (m, n),
232            "grad_c dimensions mismatch"
233        );
234        assert_eq!(b.ncols(), n, "B columns mismatch");
235        let mut grad = vec![T::default(); m.checked_mul(k).expect("matrix dimensions overflow")];
236        for j in 0..n {
237            for i in 0..m {
238                let p = self.get_argmax(i, j) as usize;
239                if p < k {
240                    grad[p * m + i] += grad_c.get_value(i, j).scalar_mul(b.get_value(p, j));
241                }
242            }
243        }
244        Mat::from_col_major(&grad, m, k)
245    }
246
247    /// MaxMul gradient with respect to B, using the winning values from A.
248    pub fn backward_b_maxmul<G>(&self, grad_c: &Mat<G>, a: &Mat<crate::TropicalMaxMul<T>>) -> Mat<G>
249    where
250        G: crate::TropicalSemiring<Scalar = T>,
251        T: Default + std::ops::AddAssign,
252    {
253        let (m, n, k) = (self.nrows(), self.ncols(), a.ncols());
254        assert_eq!(
255            (grad_c.nrows(), grad_c.ncols()),
256            (m, n),
257            "grad_c dimensions mismatch"
258        );
259        assert_eq!(a.nrows(), m, "A rows mismatch");
260        let mut grad = vec![T::default(); k.checked_mul(n).expect("matrix dimensions overflow")];
261        for j in 0..n {
262            for i in 0..m {
263                let p = self.get_argmax(i, j) as usize;
264                if p < k {
265                    grad[j * k + p] += grad_c.get_value(i, j).scalar_mul(a.get_value(i, p));
266                }
267            }
268        }
269        Mat::from_col_major(&grad, k, n)
270    }
271}
272
273#[cfg(test)]
274mod tests {
275    use super::*;
276    use crate::TropicalMaxPlus;
277
278    #[test]
279    fn test_mat_zeros() {
280        let m = Mat::<TropicalMaxPlus<f64>>::zeros(3, 4);
281        assert_eq!(m.nrows(), 3);
282        assert_eq!(m.ncols(), 4);
283        assert_eq!(m[(0, 0)].0, f64::NEG_INFINITY);
284    }
285
286    #[test]
287    fn test_mat_identity() {
288        let m = Mat::<TropicalMaxPlus<f64>>::identity(3);
289        assert_eq!(m.nrows(), 3);
290        assert_eq!(m.ncols(), 3);
291        assert_eq!(m[(0, 0)].0, 0.0); // tropical one
292        assert_eq!(m[(0, 1)].0, f64::NEG_INFINITY); // tropical zero
293        assert_eq!(m[(1, 1)].0, 0.0);
294        assert_eq!(m[(2, 2)].0, 0.0);
295    }
296
297    #[test]
298    fn test_mat_from_fn() {
299        let m =
300            Mat::<TropicalMaxPlus<f64>>::from_fn(2, 3, |i, j| TropicalMaxPlus((i * 3 + j) as f64));
301        assert_eq!(m[(0, 0)].0, 0.0);
302        assert_eq!(m[(0, 2)].0, 2.0);
303        assert_eq!(m[(1, 0)].0, 3.0);
304        assert_eq!(m[(1, 2)].0, 5.0);
305    }
306
307    #[test]
308    fn test_matref_from_slice() {
309        // Column-major data: 2×3 matrix [[1,2,3],[4,5,6]] stored as [1,4,2,5,3,6]
310        let data = [1.0f64, 4.0, 2.0, 5.0, 3.0, 6.0];
311        let m = MatRef::<TropicalMaxPlus<f64>>::from_slice(&data, 2, 3);
312        assert_eq!(m.nrows(), 2);
313        assert_eq!(m.ncols(), 3);
314        assert_eq!(m.get(0, 0), 1.0);
315        assert_eq!(m.get(1, 2), 6.0);
316    }
317
318    #[test]
319    fn test_matmul() {
320        // Column-major data:
321        // A: 2×3 matrix [[1,2,3],[4,5,6]] stored as [1,4,2,5,3,6]
322        // B: 3×2 matrix [[1,2],[3,4],[5,6]] stored as [1,3,5,2,4,6]
323        let a_data = [1.0f64, 4.0, 2.0, 5.0, 3.0, 6.0];
324        let b_data = [1.0f64, 3.0, 5.0, 2.0, 4.0, 6.0];
325
326        let a = MatRef::<TropicalMaxPlus<f64>>::from_slice(&a_data, 2, 3);
327        let b = MatRef::<TropicalMaxPlus<f64>>::from_slice(&b_data, 3, 2);
328
329        let c = a.matmul(&b);
330
331        // C[0,0] = max(1+1, 2+3, 3+5) = 8
332        assert_eq!(c[(0, 0)].0, 8.0);
333        // C[0,1] = max(1+2, 2+4, 3+6) = 9
334        assert_eq!(c[(0, 1)].0, 9.0);
335        // C[1,0] = max(4+1, 5+3, 6+5) = 11
336        assert_eq!(c[(1, 0)].0, 11.0);
337        // C[1,1] = max(4+2, 5+4, 6+6) = 12
338        assert_eq!(c[(1, 1)].0, 12.0);
339    }
340
341    #[test]
342    fn test_matmul_operator() {
343        // Column-major data
344        let a_data = [1.0f64, 4.0, 2.0, 5.0, 3.0, 6.0];
345        let b_data = [1.0f64, 3.0, 5.0, 2.0, 4.0, 6.0];
346
347        let a = MatRef::<TropicalMaxPlus<f64>>::from_slice(&a_data, 2, 3);
348        let b = MatRef::<TropicalMaxPlus<f64>>::from_slice(&b_data, 3, 2);
349
350        let c = &a * &b;
351
352        assert_eq!(c[(0, 0)].0, 8.0);
353        assert_eq!(c[(1, 1)].0, 12.0);
354    }
355
356    #[test]
357    fn test_matmul_argmax() {
358        // Column-major data
359        let a_data = [1.0f64, 4.0, 2.0, 5.0, 3.0, 6.0];
360        let b_data = [1.0f64, 3.0, 5.0, 2.0, 4.0, 6.0];
361
362        let a = MatRef::<TropicalMaxPlus<f64>>::from_slice(&a_data, 2, 3);
363        let b = MatRef::<TropicalMaxPlus<f64>>::from_slice(&b_data, 3, 2);
364
365        let result = a.matmul_argmax(&b);
366
367        assert_eq!(result.get(0, 0).0, 8.0);
368        assert_eq!(result.get_argmax(0, 0), 2); // k=2 gave max
369    }
370
371    #[test]
372    fn test_minplus_matmul() {
373        use crate::TropicalMinPlus;
374
375        // Column-major data
376        let a_data = [1.0f64, 4.0, 2.0, 5.0, 3.0, 6.0];
377        let b_data = [1.0f64, 3.0, 5.0, 2.0, 4.0, 6.0];
378
379        let a = MatRef::<TropicalMinPlus<f64>>::from_slice(&a_data, 2, 3);
380        let b = MatRef::<TropicalMinPlus<f64>>::from_slice(&b_data, 3, 2);
381
382        let c = a.matmul(&b);
383
384        // C[0,0] = min(1+1, 2+3, 3+5) = 2
385        assert_eq!(c[(0, 0)].0, 2.0);
386        // C[1,1] = min(4+2, 5+4, 6+6) = 6
387        assert_eq!(c[(1, 1)].0, 6.0);
388    }
389
390    #[test]
391    fn test_mat_as_ref() {
392        let m =
393            Mat::<TropicalMaxPlus<f64>>::from_fn(2, 3, |i, j| TropicalMaxPlus((i * 3 + j) as f64));
394
395        let r = m.as_ref();
396        assert_eq!(r.nrows(), 2);
397        assert_eq!(r.ncols(), 3);
398        assert_eq!(r.get(0, 0), 0.0);
399        assert_eq!(r.get(1, 2), 5.0);
400    }
401
402    #[test]
403    fn test_mat_matmul_direct() {
404        // Test Mat::matmul directly (no as_ref needed)
405        let a = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 4.0, 2.0, 5.0, 3.0, 6.0], 2, 3);
406        let b = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 3.0, 5.0, 2.0, 4.0, 6.0], 3, 2);
407
408        let c = a.matmul(&b);
409
410        // C[0,0] = max(1+1, 2+3, 3+5) = 8
411        assert_eq!(c[(0, 0)].0, 8.0);
412        // C[1,1] = max(4+2, 5+4, 6+6) = 12
413        assert_eq!(c[(1, 1)].0, 12.0);
414    }
415
416    #[test]
417    fn test_mat_matmul_argmax_direct() {
418        // Test Mat::matmul_argmax directly
419        let a = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 4.0, 2.0, 5.0, 3.0, 6.0], 2, 3);
420        let b = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 3.0, 5.0, 2.0, 4.0, 6.0], 3, 2);
421
422        let result = a.matmul_argmax(&b);
423
424        assert_eq!(result.get(0, 0).0, 8.0);
425        assert_eq!(result.get_argmax(0, 0), 2); // k=2 gave max
426    }
427
428    #[test]
429    fn test_mat_get_value() {
430        // Test get_value method - no trait import needed
431        let m = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 3.0, 2.0, 4.0], 2, 2);
432
433        assert_eq!(m.get_value(0, 0), 1.0);
434        assert_eq!(m.get_value(0, 1), 2.0);
435        assert_eq!(m.get_value(1, 0), 3.0);
436        assert_eq!(m.get_value(1, 1), 4.0);
437    }
438
439    #[test]
440    fn test_minplus_mat_matmul_direct() {
441        use crate::TropicalMinPlus;
442
443        let a = Mat::<TropicalMinPlus<f64>>::from_col_major(&[1.0, 4.0, 2.0, 5.0, 3.0, 6.0], 2, 3);
444        let b = Mat::<TropicalMinPlus<f64>>::from_col_major(&[1.0, 3.0, 5.0, 2.0, 4.0, 6.0], 3, 2);
445
446        let c = a.matmul(&b);
447
448        // C[0,0] = min(1+1, 2+3, 3+5) = 2
449        assert_eq!(c[(0, 0)].0, 2.0);
450        // C[1,1] = min(4+2, 5+4, 6+6) = 6
451        assert_eq!(c[(1, 1)].0, 6.0);
452    }
453
454    #[test]
455    fn test_mat_from_vec() {
456        let data = vec![
457            TropicalMaxPlus(1.0f64),
458            TropicalMaxPlus(2.0),
459            TropicalMaxPlus(3.0),
460            TropicalMaxPlus(4.0),
461        ];
462        let m = Mat::from_vec(data, 2, 2);
463        assert_eq!(m.nrows(), 2);
464        assert_eq!(m.ncols(), 2);
465        assert_eq!(m[(0, 0)].0, 1.0);
466        assert_eq!(m[(1, 1)].0, 4.0);
467    }
468
469    #[test]
470    fn test_mat_as_slice() {
471        let m = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 3.0, 2.0, 4.0], 2, 2);
472        let slice = m.as_slice();
473        assert_eq!(slice.len(), 4);
474        assert_eq!(slice[0].0, 1.0);
475        assert_eq!(slice[3].0, 4.0);
476    }
477
478    #[test]
479    fn test_mat_as_mut_slice() {
480        let mut m = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 3.0, 2.0, 4.0], 2, 2);
481        let slice = m.as_mut_slice();
482        slice[0] = TropicalMaxPlus(100.0);
483        assert_eq!(m[(0, 0)].0, 100.0);
484    }
485
486    #[test]
487    fn test_mat_as_mut_ptr() {
488        let mut m = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 3.0, 2.0, 4.0], 2, 2);
489        let ptr = m.as_mut_ptr();
490        assert!(!ptr.is_null());
491    }
492
493    #[test]
494    fn test_mat_index_mut() {
495        let mut m = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 3.0, 2.0, 4.0], 2, 2);
496        m[(0, 0)] = TropicalMaxPlus(10.0);
497        m[(1, 1)] = TropicalMaxPlus(40.0);
498        assert_eq!(m[(0, 0)].0, 10.0);
499        assert_eq!(m[(1, 1)].0, 40.0);
500    }
501
502    #[test]
503    fn test_mat_matmul_ref() {
504        let a = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 4.0, 2.0, 5.0, 3.0, 6.0], 2, 3);
505        // Column-major data for B: 3×2 matrix [[1,2],[3,4],[5,6]] stored as [1,3,5,2,4,6]
506        let b_data = [1.0f64, 3.0, 5.0, 2.0, 4.0, 6.0];
507        let b = MatRef::<TropicalMaxPlus<f64>>::from_slice(&b_data, 3, 2);
508
509        let c = a.matmul_ref(&b);
510
511        // C[0,0] = max(1+1, 2+3, 3+5) = 8
512        assert_eq!(c[(0, 0)].0, 8.0);
513        // C[1,1] = max(4+2, 5+4, 6+6) = 12
514        assert_eq!(c[(1, 1)].0, 12.0);
515    }
516
517    #[test]
518    fn test_matref_copy_clone() {
519        let data = [1.0f64, 2.0, 3.0, 4.0];
520        let a = MatRef::<TropicalMaxPlus<f64>>::from_slice(&data, 2, 2);
521        let b = a; // Copy
522        let c = a.clone(); // Clone
523        assert_eq!(a.get(0, 0), b.get(0, 0));
524        assert_eq!(a.get(0, 0), c.get(0, 0));
525    }
526
527    #[test]
528    fn test_matref_to_owned() {
529        let data = [1.0f64, 2.0, 3.0, 4.0];
530        let a = MatRef::<TropicalMaxPlus<f64>>::from_slice(&data, 2, 2);
531        let owned = a.to_owned();
532        assert_eq!(owned.nrows(), 2);
533        assert_eq!(owned.ncols(), 2);
534        assert_eq!(owned[(0, 0)].0, 1.0);
535    }
536
537    #[test]
538    fn test_matref_debug() {
539        let data = [1.0f64, 2.0];
540        let m = MatRef::<TropicalMaxPlus<f64>>::from_slice(&data, 1, 2);
541        let debug_str = format!("{:?}", m);
542        assert!(debug_str.contains("MatRef"));
543    }
544
545    #[test]
546    fn test_mat_clone() {
547        let m = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 3.0, 2.0, 4.0], 2, 2);
548        let m2 = m.clone();
549        assert_eq!(m2[(0, 0)].0, 1.0);
550        assert_eq!(m2[(1, 1)].0, 4.0);
551    }
552
553    #[test]
554    fn test_mat_debug() {
555        let m = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 2.0], 1, 2);
556        let debug_str = format!("{:?}", m);
557        assert!(debug_str.contains("Mat"));
558    }
559
560    #[test]
561    fn test_matwithargmax_get_value() {
562        let a = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 4.0, 2.0, 5.0, 3.0, 6.0], 2, 3);
563        let b = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 3.0, 5.0, 2.0, 4.0, 6.0], 3, 2);
564
565        let result = a.matmul_argmax(&b);
566
567        // Test get_value (scalar extraction without trait import)
568        assert_eq!(result.get_value(0, 0), 8.0);
569        assert_eq!(result.get_value(1, 1), 12.0);
570    }
571
572    #[test]
573    fn test_matwithargmax_nrows_ncols() {
574        let a = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 4.0, 2.0, 5.0, 3.0, 6.0], 2, 3);
575        let b = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 3.0, 5.0, 2.0, 4.0, 6.0], 3, 2);
576
577        let result = a.matmul_argmax(&b);
578
579        assert_eq!(result.nrows(), 2);
580        assert_eq!(result.ncols(), 2);
581    }
582
583    // `from_row_major` stays public (deprecated) as a downstream convenience, so keep
584    // minimal coverage of it here. All internal callers use `from_col_major` instead.
585    #[test]
586    #[should_panic(expected = "data length")]
587    #[allow(deprecated)]
588    fn test_mat_from_row_major_size_mismatch() {
589        let _ = Mat::<TropicalMaxPlus<f64>>::from_row_major(&[1.0, 2.0], 2, 2);
590    }
591
592    #[test]
593    #[allow(deprecated)]
594    fn test_from_row_major_equals_col_major_transposed() {
595        // Row-major input must yield the same logical matrix as the equivalent
596        // column-major data fed to from_col_major (this is the transpose mapping
597        // every migrated test relies on).
598        let rm = Mat::<TropicalMaxPlus<f64>>::from_row_major(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 2, 3);
599        let cm = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 4.0, 2.0, 5.0, 3.0, 6.0], 2, 3);
600        assert_eq!((rm.nrows(), rm.ncols()), (2, 3));
601        for i in 0..2 {
602            for j in 0..3 {
603                assert_eq!(rm.get_value(i, j), cm.get_value(i, j));
604            }
605        }
606    }
607
608    #[test]
609    #[should_panic(expected = "data length")]
610    fn test_mat_from_vec_size_mismatch() {
611        let data = vec![TropicalMaxPlus(1.0f64), TropicalMaxPlus(2.0)];
612        let _ = Mat::from_vec(data, 2, 2);
613    }
614
615    #[test]
616    #[should_panic(expected = "data length")]
617    fn test_matref_from_slice_size_mismatch() {
618        let data = [1.0f64, 2.0];
619        let _ = MatRef::<TropicalMaxPlus<f64>>::from_slice(&data, 2, 2);
620    }
621
622    #[test]
623    #[should_panic(expected = "dimension mismatch")]
624    fn test_matmul_dimension_mismatch() {
625        let a = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 3.0, 2.0, 4.0], 2, 2);
626        let b = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 3.0, 5.0, 2.0, 4.0, 6.0], 3, 2);
627        let _ = a.matmul(&b); // Should panic: A is 2x2, B is 3x2
628    }
629
630    #[test]
631    #[should_panic(expected = "dimension mismatch")]
632    fn test_matref_matmul_dimension_mismatch() {
633        let a_data = [1.0f64, 2.0, 3.0, 4.0];
634        let b_data = [1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0];
635        let a = MatRef::<TropicalMaxPlus<f64>>::from_slice(&a_data, 2, 2);
636        let b = MatRef::<TropicalMaxPlus<f64>>::from_slice(&b_data, 3, 2);
637        let _ = a.matmul(&b); // Should panic
638    }
639
640    #[test]
641    #[should_panic(expected = "dimension mismatch")]
642    fn test_matmul_argmax_dimension_mismatch() {
643        let a = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 3.0, 2.0, 4.0], 2, 2);
644        let b = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 3.0, 5.0, 2.0, 4.0, 6.0], 3, 2);
645        let _ = a.matmul_argmax(&b); // Should panic
646    }
647
648    #[test]
649    #[should_panic(expected = "dimension mismatch")]
650    fn test_matref_matmul_argmax_dimension_mismatch() {
651        let a_data = [1.0f64, 2.0, 3.0, 4.0];
652        let b_data = [1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0];
653        let a = MatRef::<TropicalMaxPlus<f64>>::from_slice(&a_data, 2, 2);
654        let b = MatRef::<TropicalMaxPlus<f64>>::from_slice(&b_data, 3, 2);
655        let _ = a.matmul_argmax(&b); // Should panic
656    }
657
658    #[test]
659    #[should_panic(expected = "dimension mismatch")]
660    fn test_mat_matmul_ref_dimension_mismatch() {
661        let a = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 3.0, 2.0, 4.0], 2, 2);
662        let b_data = [1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0];
663        let b = MatRef::<TropicalMaxPlus<f64>>::from_slice(&b_data, 3, 2);
664        let _ = a.matmul_ref(&b); // Should panic
665    }
666
667    // ========================================================================
668    // Batched operation tests
669    // ========================================================================
670
671    #[test]
672    fn test_mat_matmul_batched() {
673        let a1 = Mat::<TropicalMaxPlus<f32>>::from_col_major(&[1.0, 3.0, 2.0, 4.0], 2, 2);
674        let a2 = Mat::<TropicalMaxPlus<f32>>::from_col_major(&[5.0, 7.0, 6.0, 8.0], 2, 2);
675        let b1 = Mat::<TropicalMaxPlus<f32>>::from_col_major(&[1.0, 0.0, 0.0, 1.0], 2, 2);
676        let b2 = Mat::<TropicalMaxPlus<f32>>::from_col_major(&[1.0, 3.0, 2.0, 4.0], 2, 2);
677
678        let results = Mat::matmul_batched(&[a1, a2], &[b1, b2]);
679        assert_eq!(results.len(), 2);
680
681        // C[0] = A[0] * B[0] (MaxPlus)
682        // C[0,0] = max(1+1, 2+0) = 2
683        assert!((results[0][(0, 0)].0 - 2.0).abs() < 1e-5);
684
685        // C[1] = A[1] * B[1] (MaxPlus)
686        // C[0,0] = max(5+1, 6+3) = 9
687        assert!((results[1][(0, 0)].0 - 9.0).abs() < 1e-5);
688    }
689
690    #[test]
691    fn test_mat_matmul_batched_empty() {
692        let a_batch: Vec<Mat<TropicalMaxPlus<f32>>> = vec![];
693        let b_batch: Vec<Mat<TropicalMaxPlus<f32>>> = vec![];
694
695        let results = Mat::matmul_batched(&a_batch, &b_batch);
696        assert!(results.is_empty());
697    }
698
699    #[test]
700    #[should_panic(expected = "batch sizes must match")]
701    fn test_mat_matmul_batched_size_mismatch() {
702        let a1 = Mat::<TropicalMaxPlus<f32>>::from_col_major(&[1.0, 3.0, 2.0, 4.0], 2, 2);
703        let b1 = Mat::<TropicalMaxPlus<f32>>::from_col_major(&[1.0, 0.0, 0.0, 1.0], 2, 2);
704        let b2 = Mat::<TropicalMaxPlus<f32>>::from_col_major(&[1.0, 3.0, 2.0, 4.0], 2, 2);
705
706        let _ = Mat::matmul_batched(&[a1], &[b1, b2]); // Should panic
707    }
708
709    #[test]
710    #[should_panic(expected = "has dimensions")]
711    fn test_mat_matmul_batched_dimension_mismatch() {
712        let a1 = Mat::<TropicalMaxPlus<f32>>::from_col_major(&[1.0, 3.0, 2.0, 4.0], 2, 2);
713        let a2 =
714            Mat::<TropicalMaxPlus<f32>>::from_col_major(&[5.0, 8.0, 6.0, 9.0, 7.0, 10.0], 2, 3); // Different size
715        let b1 = Mat::<TropicalMaxPlus<f32>>::from_col_major(&[1.0, 0.0, 0.0, 1.0], 2, 2);
716        let b2 = Mat::<TropicalMaxPlus<f32>>::from_col_major(&[1.0, 3.0, 2.0, 4.0], 2, 2);
717
718        let _ = Mat::matmul_batched(&[a1, a2], &[b1, b2]); // Should panic
719    }
720
721    #[test]
722    fn test_mat_matmul_batched_with_argmax() {
723        let a1 = Mat::<TropicalMaxPlus<f32>>::from_col_major(&[1.0, 4.0, 2.0, 5.0, 3.0, 6.0], 2, 3);
724        let a2 = Mat::<TropicalMaxPlus<f32>>::from_col_major(&[6.0, 3.0, 5.0, 2.0, 4.0, 1.0], 2, 3);
725        let b1 = Mat::<TropicalMaxPlus<f32>>::from_col_major(&[1.0, 3.0, 5.0, 2.0, 4.0, 6.0], 3, 2);
726        let b2 = Mat::<TropicalMaxPlus<f32>>::from_col_major(&[1.0, 3.0, 5.0, 2.0, 4.0, 6.0], 3, 2);
727
728        let results = Mat::matmul_batched_with_argmax(&[a1, a2], &[b1, b2]);
729        assert_eq!(results.len(), 2);
730
731        // C[0,0] = max(1+1, 2+3, 3+5) = 8, argmax=2
732        assert!((results[0].get(0, 0).0 - 8.0).abs() < 1e-5);
733        assert_eq!(results[0].get_argmax(0, 0), 2);
734    }
735
736    #[test]
737    fn test_mat_matmul_batched_with_argmax_empty() {
738        let a_batch: Vec<Mat<TropicalMaxPlus<f32>>> = vec![];
739        let b_batch: Vec<Mat<TropicalMaxPlus<f32>>> = vec![];
740
741        let results = Mat::matmul_batched_with_argmax(&a_batch, &b_batch);
742        assert!(results.is_empty());
743    }
744
745    #[test]
746    #[should_panic(expected = "batch sizes must match")]
747    fn test_mat_matmul_batched_with_argmax_size_mismatch() {
748        let a1 = Mat::<TropicalMaxPlus<f32>>::from_col_major(&[1.0, 4.0, 2.0, 5.0, 3.0, 6.0], 2, 3);
749        let b1 = Mat::<TropicalMaxPlus<f32>>::from_col_major(&[1.0, 3.0, 5.0, 2.0, 4.0, 6.0], 3, 2);
750        let b2 = Mat::<TropicalMaxPlus<f32>>::from_col_major(&[1.0, 3.0, 5.0, 2.0, 4.0, 6.0], 3, 2);
751
752        let _ = Mat::matmul_batched_with_argmax(&[a1], &[b1, b2]); // Should panic
753    }
754
755    // ========================================================================
756    // Backward pass tests
757    // ========================================================================
758
759    #[test]
760    fn test_matwithargmax_backward_a() {
761        let a = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 4.0, 2.0, 5.0, 3.0, 6.0], 2, 3);
762        let b = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 3.0, 5.0, 2.0, 4.0, 6.0], 3, 2);
763
764        // Forward pass
765        let result = a.matmul_argmax(&b);
766
767        // All argmax should be 2 (k=2 wins for all)
768        assert_eq!(result.get_argmax(0, 0), 2);
769        assert_eq!(result.get_argmax(0, 1), 2);
770        assert_eq!(result.get_argmax(1, 0), 2);
771        assert_eq!(result.get_argmax(1, 1), 2);
772
773        // Backward pass with unit gradients
774        let grad_c = Mat::<TropicalMaxPlus<f64>>::from_fn(2, 2, |_, _| TropicalMaxPlus(1.0));
775        let grad_a = result.backward_a(&grad_c, 3);
776
777        // Only column 2 should have gradients
778        assert_eq!(grad_a.nrows(), 2);
779        assert_eq!(grad_a.ncols(), 3);
780        assert_eq!(grad_a[(0, 0)].0, 0.0); // Not selected
781        assert_eq!(grad_a[(0, 1)].0, 0.0); // Not selected
782        assert_eq!(grad_a[(0, 2)].0, 2.0); // Selected for C[0,0] and C[0,1]
783        assert_eq!(grad_a[(1, 0)].0, 0.0); // Not selected
784        assert_eq!(grad_a[(1, 1)].0, 0.0); // Not selected
785        assert_eq!(grad_a[(1, 2)].0, 2.0); // Selected for C[1,0] and C[1,1]
786    }
787
788    #[test]
789    fn test_matwithargmax_backward_b() {
790        let a = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 4.0, 2.0, 5.0, 3.0, 6.0], 2, 3);
791        let b = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 3.0, 5.0, 2.0, 4.0, 6.0], 3, 2);
792
793        // Forward pass
794        let result = a.matmul_argmax(&b);
795
796        // Backward pass with unit gradients
797        let grad_c = Mat::<TropicalMaxPlus<f64>>::from_fn(2, 2, |_, _| TropicalMaxPlus(1.0));
798        let grad_b = result.backward_b(&grad_c, 3);
799
800        // Only row 2 should have gradients
801        assert_eq!(grad_b.nrows(), 3);
802        assert_eq!(grad_b.ncols(), 2);
803        assert_eq!(grad_b[(0, 0)].0, 0.0); // Not selected
804        assert_eq!(grad_b[(0, 1)].0, 0.0); // Not selected
805        assert_eq!(grad_b[(1, 0)].0, 0.0); // Not selected
806        assert_eq!(grad_b[(1, 1)].0, 0.0); // Not selected
807        assert_eq!(grad_b[(2, 0)].0, 2.0); // Selected for C[0,0] and C[1,0]
808        assert_eq!(grad_b[(2, 1)].0, 2.0); // Selected for C[0,1] and C[1,1]
809    }
810
811    #[test]
812    fn test_matwithargmax_backward_varied_argmax() {
813        // Design matrices where different k-indices win
814        let a =
815            Mat::<TropicalMaxPlus<f64>>::from_col_major(&[10.0, 1.0, 1.0, 10.0, 1.0, 1.0], 2, 3);
816        let b =
817            Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 1.0, 10.0, 1.0, 1.0, 10.0], 3, 2);
818
819        let result = a.matmul_argmax(&b);
820
821        // Check argmax patterns
822        // C[0,0] = max(10+1=11, 1+1=2, 1+10=11), first wins -> k=0
823        // C[1,0] = max(1+1=2, 10+1=11, 1+10=11), second wins -> k=1
824        assert_eq!(result.get_argmax(0, 0), 0);
825        assert_eq!(result.get_argmax(1, 0), 1);
826
827        let grad_c = Mat::<TropicalMaxPlus<f64>>::from_fn(2, 2, |_, _| TropicalMaxPlus(1.0));
828        let grad_a = result.backward_a(&grad_c, 3);
829
830        // grad_a[0,0] should get contributions from C[0,*] where argmax == 0
831        // grad_a[1,1] should get contributions from C[1,*] where argmax == 1
832        assert!(grad_a[(0, 0)].0 > 0.0); // k=0 selected for C[0,0] and C[0,1]
833        assert!(grad_a[(1, 1)].0 > 0.0); // k=1 selected for C[1,0] and C[1,1]
834    }
835
836    #[test]
837    fn test_matwithargmax_argmax_slice() {
838        let a = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 4.0, 2.0, 5.0, 3.0, 6.0], 2, 3);
839        let b = Mat::<TropicalMaxPlus<f64>>::from_col_major(&[1.0, 3.0, 5.0, 2.0, 4.0, 6.0], 3, 2);
840
841        let result = a.matmul_argmax(&b);
842        let argmax_slice = result.argmax_slice();
843
844        assert_eq!(argmax_slice.len(), 4); // 2x2 output
845        assert_eq!(argmax_slice[0], result.get_argmax(0, 0));
846        assert_eq!(argmax_slice[1], result.get_argmax(0, 1));
847        assert_eq!(argmax_slice[2], result.get_argmax(1, 0));
848        assert_eq!(argmax_slice[3], result.get_argmax(1, 1));
849    }
850}
tropical_gemm/mat/mod.rs

tropical_gemm/mat/
mod.rs
