Struct GemmWithArgmax 

pub struct GemmWithArgmax<T: TropicalWithArgmax<Index = u32>> {
    pub values: Vec<T>,
    pub argmax: Vec<u32>,
    pub m: usize,
    pub n: usize,
    pub ld: usize,
}

Result of GEMM with argmax tracking.

For each element C[i,j], tracks which k index produced the optimal value: C[i,j] = ⊕_{k} A[i,k] ⊗ B[k,j] argmax[i,j] = argmax_k (A[i,k] ⊗ B[k,j])

Fields

values: Vec<T>

The result matrix values.

argmax: Vec<u32>

The argmax indices for each element.

m: usize

Number of rows in the result.

n: usize

Number of columns in the result.

ld: usize

Leading dimension (stride between rows).

Implementations

impl<T: TropicalWithArgmax<Index = u32>> GemmWithArgmax<T>

pub fn new(m: usize, n: usize) -> Self

Create a new result container with tropical zeros.

pub fn with_ld(m: usize, n: usize, ld: usize) -> Self

Create a new result container with specified leading dimension.

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

Get value at (i, j).

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

Get argmax at (i, j).

pub fn get_mut(&mut self, i: usize, j: usize) -> &mut T

Get mutable reference to value at (i, j).

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

Get mutable reference to argmax at (i, j).

pub fn as_mut_ptrs(&mut self) -> (*mut T, *mut u32)

Get raw pointers to the data.

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

Get the argmax array as a slice.

This is useful for backward pass computation.

pub fn values_slice(&self) -> &[T]

Get the values array as a slice.

Trait Implementations

impl<T: Clone + TropicalWithArgmax<Index = u32>> Clone for GemmWithArgmax<T>

fn clone(&self) -> GemmWithArgmax<T>

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl<T: Debug + TropicalWithArgmax<Index = u32>> Debug for GemmWithArgmax<T>

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.