tropical_gemm/types/
traits.rs

1use super::scalar::TropicalScalar;
2use std::fmt::Debug;
3
4/// Core trait for tropical semiring operations.
5///
6/// A semiring (S, ⊕, ⊗) satisfies:
7/// - (S, ⊕) is a commutative monoid with identity `tropical_zero`
8/// - (S, ⊗) is a monoid with identity `tropical_one`
9/// - ⊗ distributes over ⊕
10/// - `tropical_zero` is absorbing: a ⊗ 0 = 0 ⊗ a = 0
11pub trait TropicalSemiring: Copy + Clone + Send + Sync + Debug + PartialEq + 'static {
12    /// The underlying scalar type.
13    type Scalar: TropicalScalar;
14
15    /// Borrow scalar storage when the representation permits it.
16    /// The default uses a cached value projection in owned matrix views.
17    fn scalar_slice(_values: &[Self]) -> Option<&[Self::Scalar]> {
18        None
19    }
20
21    /// Returns the additive identity (zero element for ⊕).
22    fn tropical_zero() -> Self;
23
24    /// Returns the multiplicative identity (one element for ⊗).
25    fn tropical_one() -> Self;
26
27    /// Tropical addition (⊕).
28    fn tropical_add(self, rhs: Self) -> Self;
29
30    /// Tropical multiplication (⊗).
31    fn tropical_mul(self, rhs: Self) -> Self;
32
33    /// Get the underlying scalar value.
34    fn value(&self) -> Self::Scalar;
35
36    /// Create from a scalar value.
37    fn from_scalar(s: Self::Scalar) -> Self;
38}
39
40/// Extension trait for tropical types that support argmax tracking.
41///
42/// This is used for backpropagation: during matrix multiplication,
43/// we track which k index produced the optimal value for each C[i,j].
44pub trait TropicalWithArgmax: TropicalSemiring {
45    /// The index type used for argmax tracking.
46    type Index: Copy + Default + Debug + Send + Sync + 'static;
47
48    /// Tropical addition with argmax tracking.
49    ///
50    /// Returns the result of `tropical_add` along with the index
51    /// corresponding to which operand "won" (produced the result).
52    fn tropical_add_argmax(
53        self,
54        self_idx: Self::Index,
55        rhs: Self,
56        rhs_idx: Self::Index,
57    ) -> (Self, Self::Index);
58
59    /// Whether this (output) value is a tropical-zero "no contribution" cell
60    /// whose argmax index should be canonicalized at GEMM write-back.
61    ///
62    /// Integer tropical zeros use a guard-free in-band sentinel, so a
63    /// no-contribution cell's value drifts and its accumulated argmax adopts a
64    /// spurious `k`. Returning `true` lets the kernel reset that index to the
65    /// deterministic seed (`0`) so the whole repo agrees on one value for such
66    /// cells (and the backward pass routes no gradient there once that seed
67    /// becomes `-1`, a later step).
68    ///
69    /// Default `false`: exact-infinity types (floats) don't drift — their zero
70    /// cells already keep the seed — so the branch folds away for them.
71    #[inline(always)]
72    fn is_no_contribution(&self) -> bool {
73        false
74    }
75}
76
77/// Marker trait for tropical types that support SIMD acceleration.
78pub trait SimdTropical: TropicalSemiring {
79    /// Whether SIMD operations are available for this type.
80    const SIMD_AVAILABLE: bool;
81
82    /// The SIMD width in elements.
83    const SIMD_WIDTH: usize;
84}
85
86/// Semirings whose multiplication is ordinary addition, with unit local derivatives.
87/// Used to restrict gradient routing that does not require the input operands.
88pub trait AdditiveTropical: TropicalWithArgmax<Index = u32> {}
89impl<T: super::TropicalScalar> AdditiveTropical for super::TropicalMaxPlus<T> {}
90impl<T: super::TropicalScalar> AdditiveTropical for super::TropicalMinPlus<T> {}
tropical_gemm/types/traits.rs

tropical_gemm/types/
traits.rs
