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

tropical_gemm/mat/
mod.rs
