tropical_gemm/
api.rs

1use crate::core::{GemmWithArgmax, Transpose};
2use crate::simd::{tropical_gemm_dispatch, KernelDispatch};
3use crate::types::{TropicalSemiring, TropicalWithArgmax};
4
5#[cfg(feature = "parallel")]
6use rayon::prelude::*;
7
8fn matrix_size(rows: usize, cols: usize) -> usize {
9    rows.checked_mul(cols).expect("matrix dimensions overflow")
10}
11
12fn validate_matrix(len: usize, rows: usize, cols: usize, ld: usize, name: &str) {
13    if rows == 0 || cols == 0 {
14        return;
15    }
16    assert!(ld >= cols, "{name}: leading dimension is too small");
17    let span = (rows - 1)
18        .checked_mul(ld)
19        .and_then(|n| n.checked_add(cols))
20        .expect("matrix dimensions overflow");
21    assert!(len >= span, "{name}: buffer is too short");
22}
23
24/// Simple tropical matrix multiplication: C = A ⊗ B
25///
26/// Computes C[i,j] = ⊕_k (A[i,k] ⊗ B[k,j])
27///
28/// # Arguments
29/// - `a`: Matrix A data in row-major order
30/// - `m`: Number of rows in A
31/// - `k`: Number of columns in A / rows in B
32/// - `b`: Matrix B data in row-major order
33/// - `n`: Number of columns in B
34///
35/// # Returns
36/// Result matrix C of size m×n in row-major order
37///
38/// # Example
39///
40/// ```
41/// use tropical_gemm::{tropical_matmul, TropicalMaxPlus};
42///
43/// let a = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; // 2x3
44/// let b = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; // 3x2
45///
46/// let c = tropical_matmul::<TropicalMaxPlus<f32>>(&a, 2, 3, &b, 2);
47/// assert_eq!(c.len(), 4); // 2x2 result
48/// ```
49pub fn tropical_matmul<T: TropicalSemiring + KernelDispatch>(
50    a: &[T::Scalar],
51    m: usize,
52    k: usize,
53    b: &[T::Scalar],
54    n: usize,
55) -> Vec<T> {
56    assert_eq!(a.len(), matrix_size(m, k), "A dimensions mismatch");
57    assert_eq!(b.len(), matrix_size(k, n), "B dimensions mismatch");
58
59    let mut c = vec![T::tropical_zero(); matrix_size(m, n)];
60
61    unsafe {
62        tropical_gemm_dispatch::<T>(
63            m,
64            n,
65            k,
66            a.as_ptr(),
67            k,
68            Transpose::NoTrans,
69            b.as_ptr(),
70            n,
71            Transpose::NoTrans,
72            c.as_mut_ptr(),
73            n,
74        );
75    }
76
77    c
78}
79
80/// Tropical matrix multiplication with argmax tracking.
81///
82/// Returns both the result matrix and the argmax indices indicating
83/// which k produced each optimal C[i,j].
84///
85/// # Example
86///
87/// ```
88/// use tropical_gemm::{tropical_matmul_with_argmax, TropicalMaxPlus};
89///
90/// let a = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0]; // 2x3
91/// let b = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0]; // 3x2
92///
93/// let result = tropical_matmul_with_argmax::<TropicalMaxPlus<f64>>(&a, 2, 3, &b, 2);
94/// assert_eq!(result.m, 2);
95/// assert_eq!(result.n, 2);
96/// ```
97pub fn tropical_matmul_with_argmax<T: TropicalWithArgmax<Index = u32> + KernelDispatch>(
98    a: &[T::Scalar],
99    m: usize,
100    k: usize,
101    b: &[T::Scalar],
102    n: usize,
103) -> GemmWithArgmax<T> {
104    tropical_matmul_with_argmax_with_workspace(a, m, k, b, n, &mut crate::GemmWorkspace::new())
105}
106
107/// Tropical matrix multiplication with argmax and reusable packing storage.
108/// Output values and indices are newly allocated; workspace is reused across calls.
109pub fn tropical_matmul_with_argmax_with_workspace<
110    T: TropicalWithArgmax<Index = u32> + KernelDispatch,
111>(
112    a: &[T::Scalar],
113    m: usize,
114    k: usize,
115    b: &[T::Scalar],
116    n: usize,
117    workspace: &mut crate::GemmWorkspace<T::Scalar>,
118) -> GemmWithArgmax<T> {
119    assert_eq!(a.len(), matrix_size(m, k), "A dimensions mismatch");
120    assert_eq!(b.len(), matrix_size(k, n), "B dimensions mismatch");
121
122    let mut result = GemmWithArgmax::new(m, n);
123
124    unsafe {
125        T::dispatch_gemm_with_argmax_with_workspace(
126            m,
127            n,
128            k,
129            a.as_ptr(),
130            k,
131            Transpose::NoTrans,
132            b.as_ptr(),
133            n,
134            Transpose::NoTrans,
135            &mut result,
136            workspace,
137        );
138    }
139
140    result
141}
142
143/// Builder for configuring tropical GEMM operations.
144///
145/// Provides a fluent API for setting options like transposition,
146/// alpha/beta scaling, and output preferences.
147///
148/// # Example
149///
150/// ```
151/// use tropical_gemm::{TropicalGemm, TropicalMaxPlus, TropicalSemiring};
152///
153/// let a = vec![1.0f32; 6]; // 2x3
154/// let b = vec![1.0f32; 6]; // 3x2
155/// let mut c = vec![TropicalMaxPlus::tropical_zero(); 4]; // 2x2
156///
157/// TropicalGemm::<TropicalMaxPlus<f32>>::new(2, 2, 3)
158///     .execute(&a, 3, &b, 2, &mut c, 2);
159/// ```
160pub struct TropicalGemm<T: TropicalSemiring> {
161    m: usize,
162    n: usize,
163    k: usize,
164    trans_a: Transpose,
165    trans_b: Transpose,
166    _phantom: std::marker::PhantomData<T>,
167}
168
169impl<T: TropicalSemiring + KernelDispatch> TropicalGemm<T> {
170    /// Create a new GEMM builder.
171    pub fn new(m: usize, n: usize, k: usize) -> Self {
172        Self {
173            m,
174            n,
175            k,
176            trans_a: Transpose::NoTrans,
177            trans_b: Transpose::NoTrans,
178            _phantom: std::marker::PhantomData,
179        }
180    }
181
182    /// Transpose matrix A.
183    pub fn trans_a(mut self) -> Self {
184        self.trans_a = Transpose::Trans;
185        self
186    }
187
188    /// Transpose matrix B.
189    pub fn trans_b(mut self) -> Self {
190        self.trans_b = Transpose::Trans;
191        self
192    }
193
194    /// Execute the GEMM operation.
195    ///
196    /// # Arguments
197    /// - `a`: Matrix A data
198    /// - `lda`: Leading dimension of A
199    /// - `b`: Matrix B data
200    /// - `ldb`: Leading dimension of B
201    /// - `c`: Output matrix C (must be pre-allocated)
202    /// - `ldc`: Leading dimension of C
203    pub fn execute(
204        self,
205        a: &[T::Scalar],
206        lda: usize,
207        b: &[T::Scalar],
208        ldb: usize,
209        c: &mut [T],
210        ldc: usize,
211    ) {
212        self.execute_with_workspace(a, lda, b, ldb, c, ldc, &mut crate::GemmWorkspace::new());
213    }
214
215    /// Execute with caller-owned packing storage. Inputs and output use the
216    /// same dimensions and strides as [`Self::execute`]. Built-in semirings
217    /// reuse the buffers on subsequent calls, including parallel GEMMs.
218    #[allow(clippy::too_many_arguments)]
219    pub fn execute_with_workspace(
220        self,
221        a: &[T::Scalar],
222        lda: usize,
223        b: &[T::Scalar],
224        ldb: usize,
225        c: &mut [T],
226        ldc: usize,
227        workspace: &mut crate::GemmWorkspace<T::Scalar>,
228    ) {
229        let (ar, ac) = match self.trans_a {
230            Transpose::NoTrans => (self.m, self.k),
231            Transpose::Trans => (self.k, self.m),
232        };
233        let (br, bc) = match self.trans_b {
234            Transpose::NoTrans => (self.k, self.n),
235            Transpose::Trans => (self.n, self.k),
236        };
237        validate_matrix(a.len(), ar, ac, lda, "A");
238        validate_matrix(b.len(), br, bc, ldb, "B");
239        validate_matrix(c.len(), self.m, self.n, ldc, "C");
240        unsafe {
241            T::dispatch_gemm_with_workspace(
242                self.m,
243                self.n,
244                self.k,
245                a.as_ptr(),
246                lda,
247                self.trans_a,
248                b.as_ptr(),
249                ldb,
250                self.trans_b,
251                c.as_mut_ptr(),
252                ldc,
253                workspace,
254            );
255        }
256    }
257}
258
259/// BLAS-style GEMM interface.
260///
261/// C = A ⊗ B
262///
263/// # Safety
264/// All pointers must be valid for the specified dimensions.
265pub unsafe fn tropical_gemm<T: TropicalSemiring + KernelDispatch>(
266    m: usize,
267    n: usize,
268    k: usize,
269    a: *const T::Scalar,
270    lda: usize,
271    trans_a: Transpose,
272    b: *const T::Scalar,
273    ldb: usize,
274    trans_b: Transpose,
275    c: *mut T,
276    ldc: usize,
277) {
278    tropical_gemm_dispatch::<T>(m, n, k, a, lda, trans_a, b, ldb, trans_b, c, ldc);
279}
280
281/// Batched tropical matrix multiplication: C[i] = A[i] ⊗ B[i] for i = 0..batch_size
282///
283/// All matrices in the batch must have the same dimensions:
284/// - Each A[i] is m × k
285/// - Each B[i] is k × n
286/// - Each C[i] is m × n
287///
288/// # Arguments
289/// - `a_batch`: Slice of batch_size matrices, each of size m×k in row-major order
290/// - `b_batch`: Slice of batch_size matrices, each of size k×n in row-major order
291/// - `m`: Number of rows in each A matrix
292/// - `k`: Number of columns in A / rows in B
293/// - `n`: Number of columns in each B matrix
294///
295/// # Returns
296/// Vector of batch_size result matrices, each of size m×n
297///
298/// # Example
299///
300/// ```
301/// use tropical_gemm::{tropical_matmul_batched, TropicalMaxPlus};
302///
303/// // Two 2x2 matrix multiplications
304/// let a_batch = vec![
305///     vec![1.0f32, 2.0, 3.0, 4.0],  // A[0]: 2x2
306///     vec![5.0f32, 6.0, 7.0, 8.0],  // A[1]: 2x2
307/// ];
308/// let b_batch = vec![
309///     vec![1.0f32, 2.0, 3.0, 4.0],  // B[0]: 2x2
310///     vec![1.0f32, 2.0, 3.0, 4.0],  // B[1]: 2x2
311/// ];
312///
313/// let c_batch = tropical_matmul_batched::<TropicalMaxPlus<f32>>(&a_batch, &b_batch, 2, 2, 2);
314/// assert_eq!(c_batch.len(), 2);
315/// ```
316pub fn tropical_matmul_batched<T: TropicalSemiring + KernelDispatch>(
317    a_batch: &[Vec<T::Scalar>],
318    b_batch: &[Vec<T::Scalar>],
319    m: usize,
320    k: usize,
321    n: usize,
322) -> Vec<Vec<T>>
323where
324    T::Scalar: Send + Sync,
325    T: Send + Sync,
326{
327    assert_eq!(
328        a_batch.len(),
329        b_batch.len(),
330        "Batch sizes must match: A has {} matrices, B has {}",
331        a_batch.len(),
332        b_batch.len()
333    );
334
335    let batch_size = a_batch.len();
336    if batch_size == 0 {
337        return Vec::new();
338    }
339
340    // Validate dimensions
341    for (i, (a, b)) in a_batch.iter().zip(b_batch.iter()).enumerate() {
342        assert_eq!(
343            a.len(),
344            matrix_size(m, k),
345            "A[{}] dimensions mismatch: expected {}, got {}",
346            i,
347            matrix_size(m, k),
348            a.len()
349        );
350        assert_eq!(
351            b.len(),
352            matrix_size(k, n),
353            "B[{}] dimensions mismatch: expected {}, got {}",
354            i,
355            matrix_size(k, n),
356            b.len()
357        );
358    }
359
360    #[cfg(feature = "parallel")]
361    {
362        a_batch
363            .par_iter()
364            .zip(b_batch.par_iter())
365            .map(|(a, b)| tropical_matmul::<T>(a, m, k, b, n))
366            .collect()
367    }
368
369    #[cfg(not(feature = "parallel"))]
370    {
371        a_batch
372            .iter()
373            .zip(b_batch.iter())
374            .map(|(a, b)| tropical_matmul::<T>(a, m, k, b, n))
375            .collect()
376    }
377}
378
379/// Batched tropical matrix multiplication with argmax tracking.
380///
381/// C[i] = A[i] ⊗ B[i] for i = 0..batch_size, with argmax indices.
382///
383/// # Arguments
384/// - `a_batch`: Slice of batch_size matrices, each of size m×k
385/// - `b_batch`: Slice of batch_size matrices, each of size k×n
386/// - `m`: Number of rows in each A matrix
387/// - `k`: Number of columns in A / rows in B
388/// - `n`: Number of columns in each B matrix
389///
390/// # Returns
391/// Vector of batch_size GemmWithArgmax results
392pub fn tropical_matmul_batched_with_argmax<T: TropicalWithArgmax<Index = u32> + KernelDispatch>(
393    a_batch: &[Vec<T::Scalar>],
394    b_batch: &[Vec<T::Scalar>],
395    m: usize,
396    k: usize,
397    n: usize,
398) -> Vec<GemmWithArgmax<T>>
399where
400    T::Scalar: Send + Sync,
401    T: Send + Sync,
402{
403    assert_eq!(
404        a_batch.len(),
405        b_batch.len(),
406        "Batch sizes must match: A has {} matrices, B has {}",
407        a_batch.len(),
408        b_batch.len()
409    );
410
411    let batch_size = a_batch.len();
412    if batch_size == 0 {
413        return Vec::new();
414    }
415
416    // Validate dimensions
417    for (i, (a, b)) in a_batch.iter().zip(b_batch.iter()).enumerate() {
418        assert_eq!(
419            a.len(),
420            matrix_size(m, k),
421            "A[{}] dimensions mismatch: expected {}, got {}",
422            i,
423            matrix_size(m, k),
424            a.len()
425        );
426        assert_eq!(
427            b.len(),
428            matrix_size(k, n),
429            "B[{}] dimensions mismatch: expected {}, got {}",
430            i,
431            matrix_size(k, n),
432            b.len()
433        );
434    }
435
436    #[cfg(feature = "parallel")]
437    {
438        a_batch
439            .par_iter()
440            .zip(b_batch.par_iter())
441            .map(|(a, b)| tropical_matmul_with_argmax::<T>(a, m, k, b, n))
442            .collect()
443    }
444
445    #[cfg(not(feature = "parallel"))]
446    {
447        a_batch
448            .iter()
449            .zip(b_batch.iter())
450            .map(|(a, b)| tropical_matmul_with_argmax::<T>(a, m, k, b, n))
451            .collect()
452    }
453}
454
455/// Strided batched GEMM: computes C[i] = A[i] ⊗ B[i] from contiguous memory.
456///
457/// This is more efficient than `tropical_matmul_batched` when all matrices
458/// are stored contiguously in memory with fixed strides.
459///
460/// # Arguments
461/// - `a`: Contiguous array of all A matrices (batch_size × m × k elements)
462/// - `b`: Contiguous array of all B matrices (batch_size × k × n elements)
463/// - `batch_size`: Number of matrix pairs
464/// - `m`: Rows in each A
465/// - `k`: Columns in A / rows in B
466/// - `n`: Columns in each B
467///
468/// # Returns
469/// Contiguous array of all C matrices (batch_size × m × n elements)
470///
471/// # Example
472///
473/// ```
474/// use tropical_gemm::{tropical_matmul_strided_batched, TropicalMaxPlus};
475///
476/// // Two 2x2 matrix pairs stored contiguously
477/// let a = vec![
478///     1.0f32, 2.0, 3.0, 4.0,  // A[0]
479///     5.0, 6.0, 7.0, 8.0,      // A[1]
480/// ];
481/// let b = vec![
482///     1.0f32, 2.0, 3.0, 4.0,  // B[0]
483///     1.0, 2.0, 3.0, 4.0,      // B[1]
484/// ];
485///
486/// let c = tropical_matmul_strided_batched::<TropicalMaxPlus<f32>>(&a, &b, 2, 2, 2, 2);
487/// assert_eq!(c.len(), 8); // 2 batches × 2×2 results
488/// ```
489pub fn tropical_matmul_strided_batched<T: TropicalSemiring + KernelDispatch>(
490    a: &[T::Scalar],
491    b: &[T::Scalar],
492    batch_size: usize,
493    m: usize,
494    k: usize,
495    n: usize,
496) -> Vec<T>
497where
498    T::Scalar: Send + Sync + Copy,
499    T: Send + Sync,
500{
501    let a_stride = matrix_size(m, k);
502    let b_stride = matrix_size(k, n);
503    let c_stride = matrix_size(m, n);
504
505    assert_eq!(
506        a.len(),
507        matrix_size(batch_size, a_stride),
508        "A size mismatch: expected {}, got {}",
509        matrix_size(batch_size, a_stride),
510        a.len()
511    );
512    assert_eq!(
513        b.len(),
514        matrix_size(batch_size, b_stride),
515        "B size mismatch: expected {}, got {}",
516        matrix_size(batch_size, b_stride),
517        b.len()
518    );
519
520    if batch_size == 0 || c_stride == 0 {
521        return Vec::new();
522    }
523
524    let mut c = vec![T::tropical_zero(); matrix_size(batch_size, c_stride)];
525
526    #[cfg(feature = "parallel")]
527    {
528        c.par_chunks_mut(c_stride)
529            .enumerate()
530            .for_each(|(i, c_chunk)| {
531                let a_slice = &a[i * a_stride..(i + 1) * a_stride];
532                let b_slice = &b[i * b_stride..(i + 1) * b_stride];
533
534                unsafe {
535                    tropical_gemm_dispatch::<T>(
536                        m,
537                        n,
538                        k,
539                        a_slice.as_ptr(),
540                        k,
541                        Transpose::NoTrans,
542                        b_slice.as_ptr(),
543                        n,
544                        Transpose::NoTrans,
545                        c_chunk.as_mut_ptr(),
546                        n,
547                    );
548                }
549            });
550    }
551
552    #[cfg(not(feature = "parallel"))]
553    {
554        for i in 0..batch_size {
555            let a_slice = &a[i * a_stride..(i + 1) * a_stride];
556            let b_slice = &b[i * b_stride..(i + 1) * b_stride];
557            let c_slice = &mut c[i * c_stride..(i + 1) * c_stride];
558
559            unsafe {
560                tropical_gemm_dispatch::<T>(
561                    m,
562                    n,
563                    k,
564                    a_slice.as_ptr(),
565                    k,
566                    Transpose::NoTrans,
567                    b_slice.as_ptr(),
568                    n,
569                    Transpose::NoTrans,
570                    c_slice.as_mut_ptr(),
571                    n,
572                );
573            }
574        }
575    }
576
577    c
578}
579
580// ============================================================================
581// Backward Pass (Gradient Computation)
582// ============================================================================
583
584/// Compute gradient with respect to matrix A in tropical matmul.
585///
586/// Given the forward pass C = A ⊗ B with argmax tracking, and upstream
587/// gradient dL/dC, computes dL/dA.
588///
589/// For tropical matmul, the gradient routing is:
590/// ```text
591/// dL/dA[i,k] = Σ_j { dL/dC[i,j] if argmax[i,j] == k, else 0 }
592/// ```
593///
594/// # Arguments
595///
596/// * `grad_c` - Upstream gradient dL/dC, size m×n
597/// * `argmax` - Argmax indices from forward pass, size m×n
598/// * `m` - Number of rows in A
599/// * `k` - Number of columns in A
600/// * `n` - Number of columns in C (used for argmax indexing)
601///
602/// # Returns
603///
604/// Gradient dL/dA of size m×k
605///
606/// # Example
607///
608/// ```
609/// use tropical_gemm::{tropical_matmul_with_argmax, tropical_backward_a, TropicalMaxPlus};
610///
611/// let a = [1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0]; // 2x3
612/// let b = [1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0]; // 3x2
613///
614/// // Forward pass
615/// let result = tropical_matmul_with_argmax::<TropicalMaxPlus<f64>>(&a, 2, 3, &b, 2);
616///
617/// // Upstream gradient (e.g., all ones)
618/// let grad_c = [1.0f64; 4]; // 2x2
619///
620/// // Backward pass for A
621/// let grad_a = tropical_backward_a::<f64>(&grad_c, result.argmax_slice(), 2, 3, 2);
622/// assert_eq!(grad_a.len(), 6); // 2x3
623/// ```
624pub fn tropical_backward_a<T: Copy + Default + std::ops::AddAssign>(
625    grad_c: &[T],
626    argmax: &[u32],
627    m: usize,
628    k: usize,
629    n: usize,
630) -> Vec<T> {
631    assert_eq!(grad_c.len(), matrix_size(m, n), "grad_c size mismatch");
632    assert_eq!(argmax.len(), matrix_size(m, n), "argmax size mismatch");
633
634    let mut grad_a = vec![T::default(); matrix_size(m, k)];
635
636    for i in 0..m {
637        for j in 0..n {
638            let idx = argmax[i * n + j] as usize;
639            if idx < k {
640                grad_a[i * k + idx] += grad_c[i * n + j];
641            }
642        }
643    }
644
645    grad_a
646}
647
648/// Compute gradient with respect to matrix B in tropical matmul.
649///
650/// Given the forward pass C = A ⊗ B with argmax tracking, and upstream
651/// gradient dL/dC, computes dL/dB.
652///
653/// For tropical matmul, the gradient routing is:
654/// ```text
655/// dL/dB[k,j] = Σ_i { dL/dC[i,j] if argmax[i,j] == k, else 0 }
656/// ```
657///
658/// # Arguments
659///
660/// * `grad_c` - Upstream gradient dL/dC, size m×n
661/// * `argmax` - Argmax indices from forward pass, size m×n
662/// * `m` - Number of rows in C (used for iteration)
663/// * `k` - Number of rows in B
664/// * `n` - Number of columns in B
665///
666/// # Returns
667///
668/// Gradient dL/dB of size k×n
669///
670/// # Example
671///
672/// ```
673/// use tropical_gemm::{tropical_matmul_with_argmax, tropical_backward_b, TropicalMaxPlus};
674///
675/// let a = [1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0]; // 2x3
676/// let b = [1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0]; // 3x2
677///
678/// // Forward pass
679/// let result = tropical_matmul_with_argmax::<TropicalMaxPlus<f64>>(&a, 2, 3, &b, 2);
680///
681/// // Upstream gradient
682/// let grad_c = [1.0f64; 4]; // 2x2
683///
684/// // Backward pass for B
685/// let grad_b = tropical_backward_b::<f64>(&grad_c, result.argmax_slice(), 2, 3, 2);
686/// assert_eq!(grad_b.len(), 6); // 3x2
687/// ```
688pub fn tropical_backward_b<T: Copy + Default + std::ops::AddAssign>(
689    grad_c: &[T],
690    argmax: &[u32],
691    m: usize,
692    k: usize,
693    n: usize,
694) -> Vec<T> {
695    assert_eq!(grad_c.len(), matrix_size(m, n), "grad_c size mismatch");
696    assert_eq!(argmax.len(), matrix_size(m, n), "argmax size mismatch");
697
698    let mut grad_b = vec![T::default(); matrix_size(k, n)];
699
700    for i in 0..m {
701        for j in 0..n {
702            let idx = argmax[i * n + j] as usize;
703            if idx < k {
704                grad_b[idx * n + j] += grad_c[i * n + j];
705            }
706        }
707    }
708
709    grad_b
710}
711
712/// Batched backward pass for gradient with respect to A.
713///
714/// Computes dL/dA[i] for each batch element.
715///
716/// # Arguments
717///
718/// * `grad_c_batch` - Batch of upstream gradients, each size m×n
719/// * `argmax_batch` - Batch of argmax indices from forward pass
720/// * `m` - Number of rows in A
721/// * `k` - Number of columns in A
722/// * `n` - Number of columns in C
723///
724/// # Returns
725///
726/// Vector of gradients dL/dA[i], each of size m×k
727pub fn tropical_backward_a_batched<T: Copy + Default + std::ops::AddAssign + Send + Sync>(
728    grad_c_batch: &[Vec<T>],
729    argmax_batch: &[Vec<u32>],
730    m: usize,
731    k: usize,
732    n: usize,
733) -> Vec<Vec<T>> {
734    assert_eq!(
735        grad_c_batch.len(),
736        argmax_batch.len(),
737        "Batch sizes must match"
738    );
739
740    #[cfg(feature = "parallel")]
741    {
742        grad_c_batch
743            .par_iter()
744            .zip(argmax_batch.par_iter())
745            .map(|(grad_c, argmax)| tropical_backward_a(grad_c, argmax, m, k, n))
746            .collect()
747    }
748
749    #[cfg(not(feature = "parallel"))]
750    {
751        grad_c_batch
752            .iter()
753            .zip(argmax_batch.iter())
754            .map(|(grad_c, argmax)| tropical_backward_a(grad_c, argmax, m, k, n))
755            .collect()
756    }
757}
758
759/// Batched backward pass for gradient with respect to B.
760///
761/// Computes dL/dB[i] for each batch element.
762///
763/// # Arguments
764///
765/// * `grad_c_batch` - Batch of upstream gradients, each size m×n
766/// * `argmax_batch` - Batch of argmax indices from forward pass
767/// * `m` - Number of rows in C
768/// * `k` - Number of rows in B
769/// * `n` - Number of columns in B
770///
771/// # Returns
772///
773/// Vector of gradients dL/dB[i], each of size k×n
774pub fn tropical_backward_b_batched<T: Copy + Default + std::ops::AddAssign + Send + Sync>(
775    grad_c_batch: &[Vec<T>],
776    argmax_batch: &[Vec<u32>],
777    m: usize,
778    k: usize,
779    n: usize,
780) -> Vec<Vec<T>> {
781    assert_eq!(
782        grad_c_batch.len(),
783        argmax_batch.len(),
784        "Batch sizes must match"
785    );
786
787    #[cfg(feature = "parallel")]
788    {
789        grad_c_batch
790            .par_iter()
791            .zip(argmax_batch.par_iter())
792            .map(|(grad_c, argmax)| tropical_backward_b(grad_c, argmax, m, k, n))
793            .collect()
794    }
795
796    #[cfg(not(feature = "parallel"))]
797    {
798        grad_c_batch
799            .iter()
800            .zip(argmax_batch.iter())
801            .map(|(grad_c, argmax)| tropical_backward_b(grad_c, argmax, m, k, n))
802            .collect()
803    }
804}
805
806#[cfg(test)]
807mod tests {
808    use super::*;
809    use crate::types::TropicalMaxPlus;
810
811    #[test]
812    fn test_tropical_matmul() {
813        let a = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0];
814        let b = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0];
815
816        let c = tropical_matmul::<TropicalMaxPlus<f64>>(&a, 2, 3, &b, 2);
817
818        // C[0,0] = max(1+1, 2+3, 3+5) = 8
819        assert_eq!(c[0].0, 8.0);
820        // C[0,1] = max(1+2, 2+4, 3+6) = 9
821        assert_eq!(c[1].0, 9.0);
822        // C[1,0] = max(4+1, 5+3, 6+5) = 11
823        assert_eq!(c[2].0, 11.0);
824        // C[1,1] = max(4+2, 5+4, 6+6) = 12
825        assert_eq!(c[3].0, 12.0);
826    }
827
828    #[test]
829    fn test_tropical_matmul_with_argmax() {
830        let a = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0];
831        let b = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0];
832
833        let result = tropical_matmul_with_argmax::<TropicalMaxPlus<f64>>(&a, 2, 3, &b, 2);
834
835        assert_eq!(result.get(0, 0).0, 8.0);
836        assert_eq!(result.get_argmax(0, 0), 2); // k=2 produced max
837
838        assert_eq!(result.get(1, 1).0, 12.0);
839        assert_eq!(result.get_argmax(1, 1), 2); // k=2 produced max
840    }
841
842    #[test]
843    fn test_builder_api() {
844        let a = vec![1.0f32; 6];
845        let b = vec![1.0f32; 6];
846        let mut c = vec![TropicalMaxPlus::tropical_zero(); 4];
847
848        TropicalGemm::<TropicalMaxPlus<f32>>::new(2, 2, 3).execute(&a, 3, &b, 2, &mut c, 2);
849
850        // C[0,0] = max(1+1, 1+1, 1+1) = 2 (tropical mul is addition, tropical add is max)
851        assert_eq!(c[0].0, 2.0);
852    }
853
854    #[test]
855    fn test_builder_api_trans_a() {
856        // A is 3x2 stored as column-major (actually 2x3 in row-major transposed)
857        // A^T is 2x3, B is 3x2, result is 2x2
858        let a = vec![1.0f32, 4.0, 2.0, 5.0, 3.0, 6.0]; // col-major 3x2
859        let b = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; // row-major 3x2
860        let mut c = vec![TropicalMaxPlus::tropical_zero(); 4];
861
862        TropicalGemm::<TropicalMaxPlus<f32>>::new(2, 2, 3)
863            .trans_a()
864            .execute(&a, 2, &b, 2, &mut c, 2);
865
866        // A^T = [[1, 2, 3], [4, 5, 6]]
867        // B = [[1, 2], [3, 4], [5, 6]]
868        // C[0,0] = max(1+1, 2+3, 3+5) = 8
869        assert_eq!(c[0].0, 8.0);
870    }
871
872    #[test]
873    fn test_builder_api_trans_b() {
874        // A is 2x3, B^T is 2x3 stored as column-major, result is 2x2
875        let a = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]; // row-major 2x3
876        let b = vec![1.0f32, 3.0, 5.0, 2.0, 4.0, 6.0]; // col-major 2x3
877        let mut c = vec![TropicalMaxPlus::tropical_zero(); 4];
878
879        TropicalGemm::<TropicalMaxPlus<f32>>::new(2, 2, 3)
880            .trans_b()
881            .execute(&a, 3, &b, 3, &mut c, 2);
882
883        // A = [[1, 2, 3], [4, 5, 6]]
884        // B^T = [[1, 2], [3, 4], [5, 6]]
885        // C[0,0] = max(1+1, 2+3, 3+5) = 8
886        assert_eq!(c[0].0, 8.0);
887    }
888
889    #[test]
890    fn test_tropical_matmul_min_plus() {
891        use crate::types::TropicalMinPlus;
892
893        let a = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0];
894        let b = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0];
895
896        let c = tropical_matmul::<TropicalMinPlus<f64>>(&a, 2, 3, &b, 2);
897
898        // C[0,0] = min(1+1, 2+3, 3+5) = 2
899        assert_eq!(c[0].0, 2.0);
900        // C[0,1] = min(1+2, 2+4, 3+6) = 3
901        assert_eq!(c[1].0, 3.0);
902        // C[1,0] = min(4+1, 5+3, 6+5) = 5
903        assert_eq!(c[2].0, 5.0);
904        // C[1,1] = min(4+2, 5+4, 6+6) = 6
905        assert_eq!(c[3].0, 6.0);
906    }
907
908    #[test]
909    fn test_tropical_matmul_max_mul() {
910        use crate::types::TropicalMaxMul;
911
912        let a = vec![2.0f64, 3.0, 4.0, 5.0];
913        let b = vec![1.0f64, 2.0, 3.0, 4.0];
914
915        let c = tropical_matmul::<TropicalMaxMul<f64>>(&a, 2, 2, &b, 2);
916
917        // C[0,0] = max(2*1, 3*3) = max(2, 9) = 9
918        assert_eq!(c[0].0, 9.0);
919        // C[0,1] = max(2*2, 3*4) = max(4, 12) = 12
920        assert_eq!(c[1].0, 12.0);
921        // C[1,0] = max(4*1, 5*3) = max(4, 15) = 15
922        assert_eq!(c[2].0, 15.0);
923        // C[1,1] = max(4*2, 5*4) = max(8, 20) = 20
924        assert_eq!(c[3].0, 20.0);
925    }
926
927    #[test]
928    fn test_tropical_matmul_f32() {
929        let a = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0];
930        let b = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0];
931
932        let c = tropical_matmul::<TropicalMaxPlus<f32>>(&a, 2, 3, &b, 2);
933
934        assert!((c[0].0 - 8.0).abs() < 1e-6);
935        assert!((c[1].0 - 9.0).abs() < 1e-6);
936        assert!((c[2].0 - 11.0).abs() < 1e-6);
937        assert!((c[3].0 - 12.0).abs() < 1e-6);
938    }
939
940    #[test]
941    fn test_non_square_matrices() {
942        // 3x2 * 2x4 = 3x4
943        let a = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0];
944        let b = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
945
946        let c = tropical_matmul::<TropicalMaxPlus<f64>>(&a, 3, 2, &b, 4);
947
948        assert_eq!(c.len(), 12);
949        // C[0,0] = max(1+1, 2+5) = 7
950        assert_eq!(c[0].0, 7.0);
951    }
952
953    #[test]
954    fn test_single_element() {
955        let a = vec![5.0f64];
956        let b = vec![3.0f64];
957
958        let c = tropical_matmul::<TropicalMaxPlus<f64>>(&a, 1, 1, &b, 1);
959
960        assert_eq!(c.len(), 1);
961        assert_eq!(c[0].0, 8.0); // 5 + 3 = 8
962    }
963
964    #[test]
965    fn test_larger_matrix() {
966        let n = 16;
967        let a: Vec<f64> = (0..n * n).map(|i| i as f64).collect();
968        let b: Vec<f64> = (0..n * n).map(|i| (n * n - 1 - i) as f64).collect();
969
970        let c = tropical_matmul::<TropicalMaxPlus<f64>>(&a, n, n, &b, n);
971
972        assert_eq!(c.len(), n * n);
973        // Just verify it doesn't panic and produces reasonable results
974        for val in &c {
975            assert!(val.0.is_finite());
976        }
977    }
978
979    #[test]
980    fn test_tropical_matmul_i32() {
981        let a = vec![1i32, 2, 3, 4, 5, 6];
982        let b = vec![1i32, 2, 3, 4, 5, 6];
983
984        let c = tropical_matmul::<TropicalMaxPlus<i32>>(&a, 2, 3, &b, 2);
985
986        assert_eq!(c[0].0, 8);
987        assert_eq!(c[1].0, 9);
988        assert_eq!(c[2].0, 11);
989        assert_eq!(c[3].0, 12);
990    }
991
992    #[test]
993    fn test_tropical_matmul_i64() {
994        let a = vec![1i64, 2, 3, 4, 5, 6];
995        let b = vec![1i64, 2, 3, 4, 5, 6];
996
997        let c = tropical_matmul::<TropicalMaxPlus<i64>>(&a, 2, 3, &b, 2);
998
999        assert_eq!(c[0].0, 8);
1000        assert_eq!(c[1].0, 9);
1001        assert_eq!(c[2].0, 11);
1002        assert_eq!(c[3].0, 12);
1003    }
1004
1005    #[test]
1006    fn test_tropical_matmul_minplus_i32() {
1007        use crate::types::TropicalMinPlus;
1008
1009        let a = vec![1i32, 2, 3, 4, 5, 6];
1010        let b = vec![1i32, 2, 3, 4, 5, 6];
1011
1012        let c = tropical_matmul::<TropicalMinPlus<i32>>(&a, 2, 3, &b, 2);
1013
1014        assert_eq!(c[0].0, 2);
1015        assert_eq!(c[1].0, 3);
1016        assert_eq!(c[2].0, 5);
1017        assert_eq!(c[3].0, 6);
1018    }
1019
1020    #[test]
1021    fn test_unsafe_tropical_gemm() {
1022        let a = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0];
1023        let b = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0];
1024        let mut c = vec![TropicalMaxPlus::tropical_zero(); 4];
1025
1026        unsafe {
1027            tropical_gemm::<TropicalMaxPlus<f64>>(
1028                2,
1029                2,
1030                3,
1031                a.as_ptr(),
1032                3,
1033                Transpose::NoTrans,
1034                b.as_ptr(),
1035                2,
1036                Transpose::NoTrans,
1037                c.as_mut_ptr(),
1038                2,
1039            );
1040        }
1041
1042        assert_eq!(c[0].0, 8.0);
1043        assert_eq!(c[1].0, 9.0);
1044        assert_eq!(c[2].0, 11.0);
1045        assert_eq!(c[3].0, 12.0);
1046    }
1047
1048    #[test]
1049    fn test_minplus_with_argmax() {
1050        use crate::types::TropicalMinPlus;
1051
1052        let a = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0];
1053        let b = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0];
1054
1055        let result = tropical_matmul_with_argmax::<TropicalMinPlus<f64>>(&a, 2, 3, &b, 2);
1056
1057        // C[0,0] = min(1+1, 2+3, 3+5) = 2 at k=0
1058        assert_eq!(result.get(0, 0).0, 2.0);
1059        assert_eq!(result.get_argmax(0, 0), 0);
1060
1061        // C[1,1] = min(4+2, 5+4, 6+6) = 6 at k=0
1062        assert_eq!(result.get(1, 1).0, 6.0);
1063        assert_eq!(result.get_argmax(1, 1), 0);
1064    }
1065
1066    #[test]
1067    fn test_maxmul_with_argmax() {
1068        use crate::types::TropicalMaxMul;
1069
1070        let a = vec![2.0f64, 3.0, 4.0, 5.0];
1071        let b = vec![1.0f64, 2.0, 3.0, 4.0];
1072
1073        let result = tropical_matmul_with_argmax::<TropicalMaxMul<f64>>(&a, 2, 2, &b, 2);
1074
1075        // C[0,0] = max(2*1, 3*3) = 9 at k=1
1076        assert_eq!(result.get(0, 0).0, 9.0);
1077        assert_eq!(result.get_argmax(0, 0), 1);
1078    }
1079
1080    #[test]
1081    fn test_gemmwithargmax_dimensions() {
1082        let a = vec![1.0f64; 12]; // 3x4
1083        let b = vec![1.0f64; 20]; // 4x5
1084
1085        let result = tropical_matmul_with_argmax::<TropicalMaxPlus<f64>>(&a, 3, 4, &b, 5);
1086
1087        assert_eq!(result.m, 3);
1088        assert_eq!(result.n, 5);
1089        assert_eq!(result.values.len(), 15);
1090        assert_eq!(result.argmax.len(), 15);
1091    }
1092
1093    #[test]
1094    fn test_identity_like_matrix() {
1095        // Matrix with -inf everywhere except diagonal has 0
1096        let a = vec![0.0f64, f64::NEG_INFINITY, f64::NEG_INFINITY, 0.0];
1097        let b = vec![1.0f64, 2.0, 3.0, 4.0];
1098
1099        let c = tropical_matmul::<TropicalMaxPlus<f64>>(&a, 2, 2, &b, 2);
1100
1101        // With "identity" A, C should equal B
1102        assert_eq!(c[0].0, 1.0);
1103        assert_eq!(c[1].0, 2.0);
1104        assert_eq!(c[2].0, 3.0);
1105        assert_eq!(c[3].0, 4.0);
1106    }
1107
1108    #[test]
1109    fn test_tropical_matmul_batched() {
1110        let a_batch = vec![
1111            vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0], // 2x3
1112            vec![2.0f64, 3.0, 4.0, 5.0, 6.0, 7.0], // 2x3
1113        ];
1114        let b_batch = vec![
1115            vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0], // 3x2
1116            vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0], // 3x2
1117        ];
1118
1119        let c_batch = tropical_matmul_batched::<TropicalMaxPlus<f64>>(&a_batch, &b_batch, 2, 3, 2);
1120
1121        assert_eq!(c_batch.len(), 2);
1122
1123        // C[0][0,0] = max(1+1, 2+3, 3+5) = 8
1124        assert_eq!(c_batch[0][0].0, 8.0);
1125        // C[0][1,1] = max(4+2, 5+4, 6+6) = 12
1126        assert_eq!(c_batch[0][3].0, 12.0);
1127
1128        // C[1][0,0] = max(2+1, 3+3, 4+5) = 9
1129        assert_eq!(c_batch[1][0].0, 9.0);
1130        // C[1][1,1] = max(5+2, 6+4, 7+6) = 13
1131        assert_eq!(c_batch[1][3].0, 13.0);
1132    }
1133
1134    #[test]
1135    fn test_tropical_matmul_batched_empty() {
1136        let a_batch: Vec<Vec<f64>> = vec![];
1137        let b_batch: Vec<Vec<f64>> = vec![];
1138
1139        let c_batch = tropical_matmul_batched::<TropicalMaxPlus<f64>>(&a_batch, &b_batch, 2, 2, 2);
1140
1141        assert!(c_batch.is_empty());
1142    }
1143
1144    #[test]
1145    fn test_tropical_matmul_batched_with_argmax() {
1146        let a_batch = vec![
1147            vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0], // 2x3
1148            vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0], // 2x3
1149        ];
1150        let b_batch = vec![
1151            vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0],  // 3x2
1152            vec![10.0f64, 2.0, 3.0, 4.0, 5.0, 6.0], // 3x2 (different first element)
1153        ];
1154
1155        let results = tropical_matmul_batched_with_argmax::<TropicalMaxPlus<f64>>(
1156            &a_batch, &b_batch, 2, 3, 2,
1157        );
1158
1159        assert_eq!(results.len(), 2);
1160
1161        // First batch: C[0,0] = max(1+1, 2+3, 3+5) = 8 at k=2
1162        assert_eq!(results[0].get(0, 0).0, 8.0);
1163        assert_eq!(results[0].get_argmax(0, 0), 2);
1164
1165        // Second batch: C[0,0] = max(1+10, 2+3, 3+5) = 11 at k=0
1166        assert_eq!(results[1].get(0, 0).0, 11.0);
1167        assert_eq!(results[1].get_argmax(0, 0), 0);
1168    }
1169
1170    #[test]
1171    fn test_tropical_matmul_batched_with_argmax_empty() {
1172        let a_batch: Vec<Vec<f64>> = vec![];
1173        let b_batch: Vec<Vec<f64>> = vec![];
1174
1175        let results = tropical_matmul_batched_with_argmax::<TropicalMaxPlus<f64>>(
1176            &a_batch, &b_batch, 2, 2, 2,
1177        );
1178
1179        assert!(results.is_empty());
1180    }
1181
1182    #[test]
1183    fn test_tropical_matmul_strided_batched() {
1184        // Two 2x2 matrices stored contiguously
1185        let a = vec![
1186            1.0f64, 2.0, 3.0, 4.0, // A[0]
1187            5.0, 6.0, 7.0, 8.0, // A[1]
1188        ];
1189        let b = vec![
1190            1.0f64, 2.0, 3.0, 4.0, // B[0]
1191            1.0, 2.0, 3.0, 4.0, // B[1]
1192        ];
1193
1194        let c = tropical_matmul_strided_batched::<TropicalMaxPlus<f64>>(&a, &b, 2, 2, 2, 2);
1195
1196        assert_eq!(c.len(), 8);
1197
1198        // C[0][0,0] = max(1+1, 2+3) = 5
1199        assert_eq!(c[0].0, 5.0);
1200        // C[0][1,1] = max(3+2, 4+4) = 8
1201        assert_eq!(c[3].0, 8.0);
1202
1203        // C[1][0,0] = max(5+1, 6+3) = 9
1204        assert_eq!(c[4].0, 9.0);
1205        // C[1][1,1] = max(7+2, 8+4) = 12
1206        assert_eq!(c[7].0, 12.0);
1207    }
1208
1209    #[test]
1210    fn test_tropical_matmul_strided_batched_empty() {
1211        let a: Vec<f64> = vec![];
1212        let b: Vec<f64> = vec![];
1213
1214        let c = tropical_matmul_strided_batched::<TropicalMaxPlus<f64>>(&a, &b, 0, 2, 2, 2);
1215
1216        assert!(c.is_empty());
1217    }
1218
1219    #[test]
1220    fn test_tropical_matmul_strided_batched_minplus() {
1221        use crate::types::TropicalMinPlus;
1222
1223        let a = vec![
1224            1.0f64, 2.0, 3.0, 4.0, // A[0]
1225            5.0, 6.0, 7.0, 8.0, // A[1]
1226        ];
1227        let b = vec![
1228            1.0f64, 2.0, 3.0, 4.0, // B[0]
1229            1.0, 2.0, 3.0, 4.0, // B[1]
1230        ];
1231
1232        let c = tropical_matmul_strided_batched::<TropicalMinPlus<f64>>(&a, &b, 2, 2, 2, 2);
1233
1234        assert_eq!(c.len(), 8);
1235
1236        // C[0][0,0] = min(1+1, 2+3) = 2
1237        assert_eq!(c[0].0, 2.0);
1238        // C[0][1,1] = min(3+2, 4+4) = 5
1239        assert_eq!(c[3].0, 5.0);
1240    }
1241
1242    #[test]
1243    fn test_tropical_matmul_batched_larger() {
1244        let batch_size = 10;
1245        let m = 8;
1246        let k = 6;
1247        let n = 4;
1248
1249        let a_batch: Vec<Vec<f64>> = (0..batch_size)
1250            .map(|i| (0..m * k).map(|j| (i * m * k + j) as f64).collect())
1251            .collect();
1252        let b_batch: Vec<Vec<f64>> = (0..batch_size)
1253            .map(|_| (0..k * n).map(|j| j as f64).collect())
1254            .collect();
1255
1256        let c_batch = tropical_matmul_batched::<TropicalMaxPlus<f64>>(&a_batch, &b_batch, m, k, n);
1257
1258        assert_eq!(c_batch.len(), batch_size);
1259        for c in &c_batch {
1260            assert_eq!(c.len(), m * n);
1261            // Just verify all values are finite
1262            for val in c {
1263                assert!(val.0.is_finite());
1264            }
1265        }
1266    }
1267
1268    // ========================================================================
1269    // Backward pass tests
1270    // ========================================================================
1271
1272    #[test]
1273    fn test_tropical_backward_a() {
1274        // A is 2x3, B is 3x2, C is 2x2
1275        let a = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0];
1276        let b = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0];
1277
1278        // Forward pass
1279        let result = tropical_matmul_with_argmax::<TropicalMaxPlus<f64>>(&a, 2, 3, &b, 2);
1280
1281        // For this example:
1282        // C[0,0] = max(1+1, 2+3, 3+5) = 8, argmax=2
1283        // C[0,1] = max(1+2, 2+4, 3+6) = 9, argmax=2
1284        // C[1,0] = max(4+1, 5+3, 6+5) = 11, argmax=2
1285        // C[1,1] = max(4+2, 5+4, 6+6) = 12, argmax=2
1286        assert_eq!(result.get_argmax(0, 0), 2);
1287        assert_eq!(result.get_argmax(0, 1), 2);
1288        assert_eq!(result.get_argmax(1, 0), 2);
1289        assert_eq!(result.get_argmax(1, 1), 2);
1290
1291        // Upstream gradient (all ones)
1292        let grad_c = vec![1.0f64; 4];
1293
1294        // Backward for A
1295        let grad_a = tropical_backward_a(&grad_c, result.argmax_slice(), 2, 3, 2);
1296
1297        // Since all argmax = 2, gradients should flow to A[i,2]:
1298        // grad_a[0,0] = 0, grad_a[0,1] = 0, grad_a[0,2] = 2 (from C[0,0] and C[0,1])
1299        // grad_a[1,0] = 0, grad_a[1,1] = 0, grad_a[1,2] = 2 (from C[1,0] and C[1,1])
1300        assert_eq!(grad_a[0], 0.0); // A[0,0]
1301        assert_eq!(grad_a[1], 0.0); // A[0,1]
1302        assert_eq!(grad_a[2], 2.0); // A[0,2]
1303        assert_eq!(grad_a[3], 0.0); // A[1,0]
1304        assert_eq!(grad_a[4], 0.0); // A[1,1]
1305        assert_eq!(grad_a[5], 2.0); // A[1,2]
1306    }
1307
1308    #[test]
1309    fn test_tropical_backward_b() {
1310        let a = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0];
1311        let b = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0];
1312
1313        let result = tropical_matmul_with_argmax::<TropicalMaxPlus<f64>>(&a, 2, 3, &b, 2);
1314
1315        let grad_c = vec![1.0f64; 4];
1316
1317        // Backward for B
1318        let grad_b = tropical_backward_b(&grad_c, result.argmax_slice(), 2, 3, 2);
1319
1320        // Since all argmax = 2, gradients flow to B[2,j]:
1321        // grad_b[0,0] = 0, grad_b[0,1] = 0
1322        // grad_b[1,0] = 0, grad_b[1,1] = 0
1323        // grad_b[2,0] = 2 (from C[0,0] and C[1,0]), grad_b[2,1] = 2 (from C[0,1] and C[1,1])
1324        assert_eq!(grad_b[0], 0.0); // B[0,0]
1325        assert_eq!(grad_b[1], 0.0); // B[0,1]
1326        assert_eq!(grad_b[2], 0.0); // B[1,0]
1327        assert_eq!(grad_b[3], 0.0); // B[1,1]
1328        assert_eq!(grad_b[4], 2.0); // B[2,0]
1329        assert_eq!(grad_b[5], 2.0); // B[2,1]
1330    }
1331
1332    #[test]
1333    fn test_tropical_backward_varied_argmax() {
1334        // Design matrices where different k-indices win
1335        // A = [[10, 1], [1, 10]]
1336        // B = [[1, 10], [10, 1]]
1337        let a = vec![10.0f64, 1.0, 1.0, 10.0];
1338        let b = vec![1.0f64, 10.0, 10.0, 1.0];
1339
1340        let result = tropical_matmul_with_argmax::<TropicalMaxPlus<f64>>(&a, 2, 2, &b, 2);
1341
1342        // C[0,0] = max(10+1, 1+10) = 11, argmax=0 or 1 (tie, left wins) -> 0
1343        // C[0,1] = max(10+10, 1+1) = 20, argmax=0
1344        // C[1,0] = max(1+1, 10+10) = 20, argmax=1
1345        // C[1,1] = max(1+10, 10+1) = 11, argmax=0 or 1 (tie) -> 0
1346
1347        let grad_c = vec![1.0f64; 4];
1348        let grad_a = tropical_backward_a(&grad_c, result.argmax_slice(), 2, 2, 2);
1349        let grad_b = tropical_backward_b(&grad_c, result.argmax_slice(), 2, 2, 2);
1350
1351        // Verify gradients are distributed according to argmax
1352        assert_eq!(grad_a.len(), 4);
1353        assert_eq!(grad_b.len(), 4);
1354
1355        // The total gradient should equal the number of output elements
1356        let total_grad_a: f64 = grad_a.iter().sum();
1357        let total_grad_b: f64 = grad_b.iter().sum();
1358        assert_eq!(total_grad_a, 4.0);
1359        assert_eq!(total_grad_b, 4.0);
1360    }
1361
1362    #[test]
1363    fn test_tropical_backward_batched() {
1364        let a = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0];
1365        let b = vec![1.0f64, 2.0, 3.0, 4.0, 5.0, 6.0];
1366
1367        let result = tropical_matmul_with_argmax::<TropicalMaxPlus<f64>>(&a, 2, 3, &b, 2);
1368
1369        // Create batch
1370        let grad_c_batch = vec![vec![1.0f64; 4], vec![2.0f64; 4]];
1371        let argmax_batch = vec![
1372            result.argmax_slice().to_vec(),
1373            result.argmax_slice().to_vec(),
1374        ];
1375
1376        let grad_a_batch = tropical_backward_a_batched(&grad_c_batch, &argmax_batch, 2, 3, 2);
1377        let grad_b_batch = tropical_backward_b_batched(&grad_c_batch, &argmax_batch, 2, 3, 2);
1378
1379        assert_eq!(grad_a_batch.len(), 2);
1380        assert_eq!(grad_b_batch.len(), 2);
1381
1382        // First batch has upstream grad = 1
1383        assert_eq!(grad_a_batch[0][2], 2.0);
1384        assert_eq!(grad_b_batch[0][4], 2.0);
1385
1386        // Second batch has upstream grad = 2, so gradients should be doubled
1387        assert_eq!(grad_a_batch[1][2], 4.0);
1388        assert_eq!(grad_b_batch[1][4], 4.0);
1389    }
1390}
tropical_gemm/api.rs

tropical_gemm/
api.rs
