tropical_gemm/types/
max_mul.rs1use super::scalar::TropicalScalar;
2use super::traits::{SimdTropical, TropicalSemiring, TropicalWithArgmax};
3use std::fmt;
4use std::ops::{Add, Mul};
5
6#[derive(Copy, Clone, PartialEq)]
17#[repr(transparent)]
18pub struct TropicalMaxMul<T: TropicalScalar>(pub T);
19
20impl<T: TropicalScalar> TropicalMaxMul<T> {
21 #[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 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 assert_eq!(a.tropical_add(zero), a);
141 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 assert_eq!(a.tropical_add(b).0, 5.0);
152 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 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 assert_eq!((a + b).0, 5.0);
172 assert_eq!((b + a).0, 5.0);
173
174 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); 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 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]; 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 let high = TropicalMaxMul::new(0.9f64);
297 let medium = TropicalMaxMul::new(0.5f64);
298 let low = TropicalMaxMul::new(0.2f64);
299
300 let or_result = high.tropical_add(low);
302 assert_eq!(or_result.0, 0.9);
303
304 let and_result = high.tropical_mul(medium);
306 assert!((and_result.0 - 0.45).abs() < 1e-10);
307 }
308}