Struct MatWithArgmax 

pub struct MatWithArgmax<S: TropicalWithArgmax> {
    pub values: Mat<S>,
    pub argmax: Vec<u32>,
}

Result of matrix multiplication with argmax tracking.

Fields

values: Mat<S>

The result matrix values.

argmax: Vec<u32>

The argmax indices (which k produced each C[i,j]).

Implementations

impl<S: TropicalWithArgmax<Index = u32>> MatWithArgmax<S>

pub fn get(&self, i: usize, j: usize) -> S

Get the value at position (i, j).

pub fn get_value(&self, i: usize, j: usize) -> S::Scalar

Get the scalar value at position (i, j).

This is a convenience method that extracts the underlying scalar without requiring a trait import.

pub fn get_argmax(&self, i: usize, j: usize) -> u32

Get the argmax index at position (i, j).

pub fn nrows(&self) -> usize

Number of rows.

pub fn ncols(&self) -> usize

Number of columns.

pub fn argmax_slice(&self) -> &[u32]

Get the argmax indices as a slice.

This is useful for backward pass computation.

pub fn backward_a<G>(&self, grad_c: &Mat<G>, k: usize) -> Mat<G>

where G: TropicalSemiring, G::Scalar: Copy + Default + AddAssign,

Compute gradient with respect to matrix A.

Given the upstream gradient dL/dC, computes dL/dA using the argmax indices from the forward pass.

For C = A ⊗ B where C[i,j] = ⊕_k (A[i,k] ⊗ B[k,j]): dL/dA[i,k] = Σ_j { dL/dC[i,j] if argmax[i,j] == k }

Arguments

• grad_c - Gradient of the loss with respect to C, dimensions m×n
• k - Number of columns in A (the inner dimension)

Returns

Gradient of the loss with respect to A, dimensions m×k

Example

use tropical_gemm::{Mat, MaxPlus, TropicalMaxPlus};

let a = Mat::<MaxPlus<f64>>::from_row_major(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 2, 3);
let b = Mat::<MaxPlus<f64>>::from_row_major(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 3, 2);

// Forward pass with argmax
let result = a.matmul_argmax(&b);

// Backward pass: grad_c is upstream gradient (e.g., all ones)
let grad_c = Mat::<MaxPlus<f64>>::from_fn(2, 2, |_, _| TropicalMaxPlus(1.0));
let grad_a = result.backward_a(&grad_c, 3); // k=3 (columns in A)

assert_eq!(grad_a.nrows(), 2);
assert_eq!(grad_a.ncols(), 3);

pub fn backward_b<G>(&self, grad_c: &Mat<G>, k: usize) -> Mat<G>

where G: TropicalSemiring, G::Scalar: Copy + Default + AddAssign,

Compute gradient with respect to matrix B.

Given the upstream gradient dL/dC, computes dL/dB using the argmax indices from the forward pass.

For C = A ⊗ B where C[i,j] = ⊕_k (A[i,k] ⊗ B[k,j]): dL/dB[k,j] = Σ_i { dL/dC[i,j] if argmax[i,j] == k }

Arguments

• grad_c - Gradient of the loss with respect to C, dimensions m×n
• k - Number of rows in B (the inner dimension)

Returns

Gradient of the loss with respect to B, dimensions k×n

Example

use tropical_gemm::{Mat, MaxPlus, TropicalMaxPlus};

let a = Mat::<MaxPlus<f64>>::from_row_major(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 2, 3);
let b = Mat::<MaxPlus<f64>>::from_row_major(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], 3, 2);

// Forward pass with argmax
let result = a.matmul_argmax(&b);

// Backward pass: grad_c is upstream gradient
let grad_c = Mat::<MaxPlus<f64>>::from_fn(2, 2, |_, _| TropicalMaxPlus(1.0));
let grad_b = result.backward_b(&grad_c, 3); // k=3 (rows in B)

assert_eq!(grad_b.nrows(), 3);
assert_eq!(grad_b.ncols(), 2);