Skip to main content

TropicalWithArgmax

tropical_gemm::types

Trait TropicalWithArgmax 

Source 
pub trait TropicalWithArgmax: TropicalSemiring {
    type Index: Copy + Default + Debug + Send + Sync + 'static;

    // Required method
    fn tropical_add_argmax(
        self,
        self_idx: Self::Index,
        rhs: Self,
        rhs_idx: Self::Index,
    ) -> (Self, Self::Index);

    // Provided method
    fn is_no_contribution(&self) -> bool { ... }
}
Expand description

Extension trait for tropical types that support argmax tracking.

This is used for backpropagation: during matrix multiplication, we track which k index produced the optimal value for each C[i,j].

Required Associated Types§

Source

type Index: Copy + Default + Debug + Send + Sync + 'static

The index type used for argmax tracking.

Required Methods§

Source

fn tropical_add_argmax( self, self_idx: Self::Index, rhs: Self, rhs_idx: Self::Index, ) -> (Self, Self::Index)

Tropical addition with argmax tracking.

Returns the result of tropical_add along with the index corresponding to which operand “won” (produced the result).

Provided Methods§

Source

fn is_no_contribution(&self) -> bool

Whether this (output) value is a tropical-zero “no contribution” cell whose argmax index should be canonicalized at GEMM write-back.

Integer tropical zeros use a guard-free in-band sentinel, so a no-contribution cell’s value drifts and its accumulated argmax adopts a spurious k. Returning true lets the kernel reset that index to the deterministic seed (0) so the whole repo agrees on one value for such cells (and the backward pass routes no gradient there once that seed becomes -1, a later step).

Default false: exact-infinity types (floats) don’t drift — their zero cells already keep the seed — so the branch folds away for them.

Dyn Compatibility§

This trait is not dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety", so this trait is not object safe.

Implementors§