Skip to main content

tropical_gemm/types/
max_mul.rs

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