tropical_gemm/types/
max_plus.rs

1use super::scalar::TropicalScalar;
2use super::traits::{SimdTropical, TropicalSemiring, TropicalWithArgmax};
3use std::fmt;
4use std::ops::{Add, Mul};
5
6/// TropicalMaxPlus semiring: (ℝ ∪ {-∞}, max, +)
7///
8/// - Addition (⊕) = max
9/// - Multiplication (⊗) = +
10/// - Zero = -∞
11/// - One = 0
12///
13/// This is the classic tropical semiring used in:
14/// - Viterbi algorithm
15/// - Shortest path algorithms (with negated weights)
16/// - Log-space probability computations
17#[derive(Copy, Clone, PartialEq)]
18#[repr(transparent)]
19pub struct TropicalMaxPlus<T: TropicalScalar>(pub T);
20
21impl<T: TropicalScalar> TropicalMaxPlus<T> {
22    /// Create a new TropicalMaxPlus value.
23    #[inline(always)]
24    pub fn new(value: T) -> Self {
25        Self(value)
26    }
27}
28
29impl<T: TropicalScalar> TropicalSemiring for TropicalMaxPlus<T> {
30    type Scalar = T;
31
32    fn scalar_slice(values: &[Self]) -> Option<&[Self::Scalar]> {
33        // SAFETY: this type is repr(transparent) over its scalar field.
34        Some(unsafe { std::slice::from_raw_parts(values.as_ptr().cast(), values.len()) })
35    }
36
37    #[inline(always)]
38    fn tropical_zero() -> Self {
39        Self(T::neg_infinity())
40    }
41
42    #[inline(always)]
43    fn tropical_one() -> Self {
44        Self(T::scalar_zero())
45    }
46
47    #[inline(always)]
48    fn tropical_add(self, rhs: Self) -> Self {
49        Self(self.0.scalar_max(rhs.0))
50    }
51
52    #[inline(always)]
53    fn tropical_mul(self, rhs: Self) -> Self {
54        Self(self.0.scalar_add(rhs.0))
55    }
56
57    #[inline(always)]
58    fn value(&self) -> T {
59        self.0
60    }
61
62    #[inline(always)]
63    fn from_scalar(s: T) -> Self {
64        Self(s)
65    }
66}
67
68impl<T: TropicalScalar> TropicalWithArgmax for TropicalMaxPlus<T> {
69    type Index = u32;
70
71    #[inline(always)]
72    fn tropical_add_argmax(self, self_idx: u32, rhs: Self, rhs_idx: u32) -> (Self, u32) {
73        if self.0 >= rhs.0 {
74            (self, self_idx)
75        } else {
76            (rhs, rhs_idx)
77        }
78    }
79
80    #[inline(always)]
81    fn is_no_contribution(&self) -> bool {
82        self.0.is_drifted_neg_zero()
83    }
84}
85
86impl<T: TropicalScalar> SimdTropical for TropicalMaxPlus<T> {
87    const SIMD_AVAILABLE: bool = true;
88    const SIMD_WIDTH: usize = 8; // f32x8 for AVX2
89}
90
91impl<T: TropicalScalar> Add for TropicalMaxPlus<T> {
92    type Output = Self;
93
94    #[inline(always)]
95    fn add(self, rhs: Self) -> Self::Output {
96        self.tropical_add(rhs)
97    }
98}
99
100impl<T: TropicalScalar> Mul for TropicalMaxPlus<T> {
101    type Output = Self;
102
103    #[inline(always)]
104    fn mul(self, rhs: Self) -> Self::Output {
105        self.tropical_mul(rhs)
106    }
107}
108
109impl<T: TropicalScalar> Default for TropicalMaxPlus<T> {
110    #[inline(always)]
111    fn default() -> Self {
112        Self::tropical_zero()
113    }
114}
115
116impl<T: TropicalScalar> fmt::Debug for TropicalMaxPlus<T> {
117    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
118        write!(f, "TropicalMaxPlus({})", self.0)
119    }
120}
121
122impl<T: TropicalScalar> fmt::Display for TropicalMaxPlus<T> {
123    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
124        write!(f, "{}", self.0)
125    }
126}
127
128impl<T: TropicalScalar> From<T> for TropicalMaxPlus<T> {
129    #[inline(always)]
130    fn from(value: T) -> Self {
131        Self(value)
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use super::*;
138
139    #[test]
140    fn test_semiring_identity() {
141        let a = TropicalMaxPlus::new(5.0f64);
142        let zero = TropicalMaxPlus::tropical_zero();
143        let one = TropicalMaxPlus::tropical_one();
144
145        // a ⊕ 0 = a
146        assert_eq!(a.tropical_add(zero), a);
147        // a ⊗ 1 = a
148        assert_eq!(a.tropical_mul(one), a);
149    }
150
151    #[test]
152    fn test_operations() {
153        let a = TropicalMaxPlus::new(3.0f64);
154        let b = TropicalMaxPlus::new(5.0f64);
155
156        // max(3, 5) = 5
157        assert_eq!(a.tropical_add(b).0, 5.0);
158        // 3 + 5 = 8
159        assert_eq!(a.tropical_mul(b).0, 8.0);
160    }
161
162    #[test]
163    fn test_argmax() {
164        let a = TropicalMaxPlus::new(3.0f64);
165        let b = TropicalMaxPlus::new(5.0f64);
166
167        let (result, idx) = a.tropical_add_argmax(0, b, 1);
168        assert_eq!(result.0, 5.0);
169        assert_eq!(idx, 1);
170    }
171
172    #[test]
173    fn test_argmax_left_wins() {
174        let a = TropicalMaxPlus::new(7.0f64);
175        let b = TropicalMaxPlus::new(3.0f64);
176
177        let (result, idx) = a.tropical_add_argmax(10, b, 20);
178        assert_eq!(result.0, 7.0);
179        assert_eq!(idx, 10); // Left wins, keep left index
180    }
181
182    #[test]
183    fn test_argmax_equal_values() {
184        // When values are equal, left (self) wins (>= comparison)
185        let a = TropicalMaxPlus::new(5.0f64);
186        let b = TropicalMaxPlus::new(5.0f64);
187
188        let (result, idx) = a.tropical_add_argmax(1, b, 2);
189        assert_eq!(result.0, 5.0);
190        assert_eq!(idx, 1); // Equal, so left (self) wins
191    }
192
193    #[test]
194    fn test_argmax_chain() {
195        // Simulate accumulating through k iterations
196        let mut acc = TropicalMaxPlus::tropical_zero();
197        let mut idx = 0u32;
198
199        let values = [3.0, 7.0, 2.0, 5.0]; // Max is at index 1
200        for (k, &val) in values.iter().enumerate() {
201            let candidate = TropicalMaxPlus::new(val);
202            (acc, idx) = acc.tropical_add_argmax(idx, candidate, k as u32);
203        }
204
205        assert_eq!(acc.0, 7.0);
206        assert_eq!(idx, 1); // Index where max occurred
207    }
208
209    #[test]
210    fn test_argmax_neg_infinity() {
211        let a = TropicalMaxPlus::tropical_zero(); // -inf
212        let b = TropicalMaxPlus::new(-100.0f64);
213
214        let (result, idx) = a.tropical_add_argmax(0, b, 1);
215        assert_eq!(result.0, -100.0);
216        assert_eq!(idx, 1); // -100 > -inf
217    }
218
219    #[test]
220    fn test_absorbing_zero() {
221        let a = TropicalMaxPlus::new(5.0f64);
222        let zero = TropicalMaxPlus::tropical_zero();
223
224        // a ⊗ 0 = a + (-inf) = -inf
225        // In tropical max-plus, multiplying by zero (adding -inf) gives -inf
226        let result = a.tropical_mul(zero);
227        assert!(result.0.is_infinite() && result.0 < 0.0);
228    }
229
230    #[test]
231    fn test_operator_overloads() {
232        let a = TropicalMaxPlus::new(3.0f64);
233        let b = TropicalMaxPlus::new(5.0f64);
234
235        // Add operator (max)
236        assert_eq!((a + b).0, 5.0);
237        assert_eq!((b + a).0, 5.0);
238
239        // Mul operator (add)
240        assert_eq!((a * b).0, 8.0);
241        assert_eq!((b * a).0, 8.0);
242    }
243
244    #[test]
245    fn test_default() {
246        let d = TropicalMaxPlus::<f64>::default();
247        assert!(d.0.is_infinite() && d.0 < 0.0); // -inf
248        assert_eq!(d, TropicalMaxPlus::tropical_zero());
249    }
250
251    #[test]
252    fn test_display_debug() {
253        let a = TropicalMaxPlus::new(5.0f64);
254
255        assert_eq!(format!("{}", a), "5");
256        assert_eq!(format!("{:?}", a), "TropicalMaxPlus(5)");
257    }
258
259    #[test]
260    fn test_from() {
261        let a: TropicalMaxPlus<f64> = 5.0.into();
262        assert_eq!(a.0, 5.0);
263
264        let b = TropicalMaxPlus::<f64>::from(3.0);
265        assert_eq!(b.0, 3.0);
266    }
267
268    #[test]
269    fn test_value_and_from_scalar() {
270        let a = TropicalMaxPlus::new(5.0f64);
271        assert_eq!(a.value(), 5.0);
272
273        let b = TropicalMaxPlus::<f64>::from_scalar(3.0);
274        assert_eq!(b.value(), 3.0);
275    }
276
277    #[test]
278    fn test_simd_tropical() {
279        assert!(TropicalMaxPlus::<f64>::SIMD_AVAILABLE);
280        assert_eq!(TropicalMaxPlus::<f64>::SIMD_WIDTH, 8);
281    }
282
283    #[test]
284    fn test_clone_copy() {
285        let a = TropicalMaxPlus::new(5.0f64);
286        let a_copy = a;
287        let a_clone = a.clone();
288
289        assert_eq!(a, a_copy);
290        assert_eq!(a, a_clone);
291    }
292
293    #[test]
294    fn test_eq() {
295        let a1 = TropicalMaxPlus::new(5.0f64);
296        let a2 = TropicalMaxPlus::new(5.0f64);
297        let b = TropicalMaxPlus::new(3.0f64);
298
299        assert_eq!(a1, a2);
300        assert_ne!(a1, b);
301    }
302
303    #[test]
304    fn test_f32() {
305        let a = TropicalMaxPlus::new(3.0f32);
306        let b = TropicalMaxPlus::new(5.0f32);
307
308        assert!((a.tropical_add(b).0 - 5.0).abs() < 1e-6);
309        assert!((a.tropical_mul(b).0 - 8.0).abs() < 1e-6);
310    }
311}
tropical_gemm/types/max_plus.rs

tropical_gemm/types/
max_plus.rs
