rstsr_core/device_cpu_serial/
reduction.rs

1use crate::prelude_dev::*;
2use core::ops::{Add, Mul};
3use num::complex::ComplexFloat;
4use num::{FromPrimitive, One, Zero};
5use rstsr_dtype_traits::ExtReal;
6
7impl<T, D> OpSumAPI<T, D> for DeviceCpuSerial
8where
9    T: Zero + Add<Output = T> + Clone,
10    D: DimAPI,
11{
12    type TOut = T;
13
14    fn sum_all(&self, a: &Vec<T>, la: &Layout<D>) -> Result<T> {
15        let f_init = T::zero;
16        let f = |acc, x| acc + x;
17        let f_sum = |acc1, acc2| acc1 + acc2;
18        let f_out = |acc| acc;
19
20        reduce_all_cpu_serial(a, la, f_init, f, f_sum, f_out)
21    }
22
23    fn sum_axes(
24        &self,
25        a: &Vec<T>,
26        la: &Layout<D>,
27        axes: &[isize],
28    ) -> Result<(Storage<DataOwned<Vec<T>>, T, Self>, Layout<IxD>)> {
29        let f_init = T::zero;
30        let f = |acc, x| acc + x;
31        let f_sum = |acc1, acc2| acc1 + acc2;
32        let f_out = |acc| acc;
33
34        let (out, layout_out) = reduce_axes_cpu_serial(a, &la.to_dim()?, axes, f_init, f, f_sum, f_out)?;
35        Ok((Storage::new(out.into(), self.clone()), layout_out))
36    }
37}
38
39impl<T, D> OpMinAPI<T, D> for DeviceCpuSerial
40where
41    T: ExtReal,
42    D: DimAPI,
43{
44    type TOut = T;
45
46    fn min_all(&self, a: &Vec<T>, la: &Layout<D>) -> Result<T> {
47        if la.size() == 0 {
48            rstsr_raise!(InvalidValue, "zero-size array is not supported for min")?;
49        }
50
51        let f_init = T::ext_max_value;
52        let f = |acc: T, x: T| acc.ext_min(x);
53        let f_sum = |acc1: T, acc2: T| acc1.ext_min(acc2);
54        let f_out = |acc| acc;
55        reduce_all_cpu_serial(a, la, f_init, f, f_sum, f_out)
56    }
57
58    fn min_axes(
59        &self,
60        a: &Vec<T>,
61        la: &Layout<D>,
62        axes: &[isize],
63    ) -> Result<(Storage<DataOwned<Vec<T>>, T, Self>, Layout<IxD>)> {
64        if la.size() == 0 {
65            rstsr_raise!(InvalidValue, "zero-size array is not supported for min")?;
66        }
67
68        let f_init = T::ext_max_value;
69        let f = |acc: T, x: T| acc.ext_min(x);
70        let f_sum = |acc1: T, acc2: T| acc1.ext_min(acc2);
71        let f_out = |acc| acc;
72
73        let (out, layout_out) = reduce_axes_cpu_serial(a, &la.to_dim()?, axes, f_init, f, f_sum, f_out)?;
74        Ok((Storage::new(out.into(), self.clone()), layout_out))
75    }
76}
77
78impl<T, D> OpMaxAPI<T, D> for DeviceCpuSerial
79where
80    T: ExtReal,
81    D: DimAPI,
82{
83    type TOut = T;
84
85    fn max_all(&self, a: &Vec<T>, la: &Layout<D>) -> Result<T> {
86        if la.size() == 0 {
87            rstsr_raise!(InvalidValue, "zero-size array is not supported for max")?;
88        }
89
90        let f_init = T::ext_min_value;
91        let f = |acc: T, x: T| acc.ext_max(x);
92        let f_sum = |acc1: T, acc2: T| acc1.ext_max(acc2);
93        let f_out = |acc| acc;
94
95        reduce_all_cpu_serial(a, la, f_init, f, f_sum, f_out)
96    }
97
98    fn max_axes(
99        &self,
100        a: &Vec<T>,
101        la: &Layout<D>,
102        axes: &[isize],
103    ) -> Result<(Storage<DataOwned<Vec<T>>, T, Self>, Layout<IxD>)> {
104        if la.size() == 0 {
105            rstsr_raise!(InvalidValue, "zero-size array is not supported for max")?;
106        }
107
108        let f_init = T::ext_min_value;
109        let f = |acc: T, x: T| acc.ext_max(x);
110        let f_sum = |acc1: T, acc2: T| acc1.ext_max(acc2);
111        let f_out = |acc| acc;
112
113        let (out, layout_out) = reduce_axes_cpu_serial(a, &la.to_dim()?, axes, f_init, f, f_sum, f_out)?;
114        Ok((Storage::new(out.into(), self.clone()), layout_out))
115    }
116}
117
118impl<T, D> OpProdAPI<T, D> for DeviceCpuSerial
119where
120    T: Clone + One + Mul<Output = T>,
121    D: DimAPI,
122{
123    type TOut = T;
124
125    fn prod_all(&self, a: &Vec<T>, la: &Layout<D>) -> Result<T> {
126        let f_init = T::one;
127        let f = |acc, x| acc * x;
128        let f_sum = |acc1, acc2| acc1 * acc2;
129        let f_out = |acc| acc;
130
131        reduce_all_cpu_serial(a, la, f_init, f, f_sum, f_out)
132    }
133
134    fn prod_axes(
135        &self,
136        a: &Vec<T>,
137        la: &Layout<D>,
138        axes: &[isize],
139    ) -> Result<(Storage<DataOwned<Vec<T>>, T, Self>, Layout<IxD>)> {
140        let f_init = T::one;
141        let f = |acc, x| acc * x;
142        let f_sum = |acc1, acc2| acc1 * acc2;
143        let f_out = |acc| acc;
144
145        let (out, layout_out) = reduce_axes_cpu_serial(a, &la.to_dim()?, axes, f_init, f, f_sum, f_out)?;
146        Ok((Storage::new(out.into(), self.clone()), layout_out))
147    }
148}
149
150impl<T, D> OpMeanAPI<T, D> for DeviceCpuSerial
151where
152    T: Clone + ComplexFloat + FromPrimitive,
153    D: DimAPI,
154{
155    type TOut = T;
156
157    fn mean_all(&self, a: &Vec<T>, la: &Layout<D>) -> Result<T> {
158        let size = la.size();
159        let f_init = T::zero;
160        let f = |acc, x| acc + x;
161        let f_sum = |acc, x| acc + x;
162        let f_out = |acc| acc / T::from_usize(size).unwrap();
163
164        let sum = reduce_all_cpu_serial(a, la, f_init, f, f_sum, f_out)?;
165        Ok(sum)
166    }
167
168    fn mean_axes(
169        &self,
170        a: &Vec<T>,
171        la: &Layout<D>,
172        axes: &[isize],
173    ) -> Result<(Storage<DataOwned<Vec<T>>, T, Self>, Layout<IxD>)> {
174        let (layout_axes, _) = la.dim_split_axes(axes)?;
175        let size = layout_axes.size();
176        let f_init = T::zero;
177        let f = |acc, x| acc + x;
178        let f_sum = |acc, x| acc + x;
179        let f_out = |acc| acc / T::from_usize(size).unwrap();
180
181        let (out, layout_out) = reduce_axes_cpu_serial(a, &la.to_dim()?, axes, f_init, f, f_sum, f_out)?;
182        Ok((Storage::new(out.into(), self.clone()), layout_out))
183    }
184}
185
186impl<T, D> OpVarAPI<T, D> for DeviceCpuSerial
187where
188    T: Clone + ComplexFloat + FromPrimitive,
189    T::Real: Clone + ComplexFloat + FromPrimitive,
190    D: DimAPI,
191{
192    type TOut = T::Real;
193
194    fn var_all(&self, a: &Vec<T>, la: &Layout<D>) -> Result<T::Real> {
195        let size = la.size();
196
197        let f_init = || (T::zero(), T::Real::zero());
198        let f = |(acc_1, acc_2): (T, T::Real), x: T| (acc_1 + x, acc_2 + (x * x.conj()).re());
199        let f_sum = |(acc_1, acc_2): (T, T::Real), (x_1, x_2)| (acc_1 + x_1, acc_2 + x_2);
200        let f_out = |(acc_1, acc_2): (T, T::Real)| {
201            let size_1 = T::from_usize(size).unwrap();
202            let size_2 = T::Real::from_usize(size).unwrap();
203            let mean = acc_1 / size_1;
204            acc_2 / size_2 - (mean * mean.conj()).re()
205        };
206
207        let result = reduce_all_cpu_serial(a, la, f_init, f, f_sum, f_out)?;
208        Ok(result)
209    }
210
211    fn var_axes(
212        &self,
213        a: &Vec<T>,
214        la: &Layout<D>,
215        axes: &[isize],
216    ) -> Result<(Storage<DataOwned<Vec<T::Real>>, T::Real, Self>, Layout<IxD>)> {
217        let (layout_axes, _) = la.dim_split_axes(axes)?;
218        let size = layout_axes.size();
219
220        let f_init = || (T::zero(), T::Real::zero());
221        let f = |(acc_1, acc_2): (T, T::Real), x: T| (acc_1 + x, acc_2 + (x * x.conj()).re());
222        let f_sum = |(acc_1, acc_2): (T, T::Real), (x_1, x_2)| (acc_1 + x_1, acc_2 + x_2);
223        let f_out = |(acc_1, acc_2): (T, T::Real)| {
224            let size_1 = T::from_usize(size).unwrap();
225            let size_2 = T::Real::from_usize(size).unwrap();
226            let mean = acc_1 / size_1;
227            acc_2 / size_2 - (mean * mean.conj()).re()
228        };
229
230        let (out, layout_out) = reduce_axes_difftype_cpu_serial(a, &la.to_dim()?, axes, f_init, f, f_sum, f_out)?;
231
232        Ok((Storage::new(out.into(), self.clone()), layout_out))
233    }
234}
235
236impl<T, D> OpStdAPI<T, D> for DeviceCpuSerial
237where
238    T: Clone + ComplexFloat + FromPrimitive,
239    T::Real: Clone + ComplexFloat + FromPrimitive,
240    D: DimAPI,
241{
242    type TOut = T::Real;
243
244    fn std_all(&self, a: &Vec<T>, la: &Layout<D>) -> Result<T::Real> {
245        let size = la.size();
246
247        let f_init = || (T::zero(), T::Real::zero());
248        let f = |(acc_1, acc_2): (T, T::Real), x: T| (acc_1 + x, acc_2 + (x * x.conj()).re());
249        let f_sum = |(acc_1, acc_2): (T, T::Real), (x_1, x_2)| (acc_1 + x_1, acc_2 + x_2);
250        let f_out = |(acc_1, acc_2): (T, T::Real)| {
251            let size_1 = T::from_usize(size).unwrap();
252            let size_2 = T::Real::from_usize(size).unwrap();
253            let mean = acc_1 / size_1;
254            let var = acc_2 / size_2 - (mean * mean.conj()).re();
255            var.sqrt()
256        };
257
258        let result = reduce_all_cpu_serial(a, la, f_init, f, f_sum, f_out)?;
259        Ok(result)
260    }
261
262    fn std_axes(
263        &self,
264        a: &Vec<T>,
265        la: &Layout<D>,
266        axes: &[isize],
267    ) -> Result<(Storage<DataOwned<Vec<T::Real>>, T::Real, Self>, Layout<IxD>)> {
268        let (layout_axes, _) = la.dim_split_axes(axes)?;
269        let size = layout_axes.size();
270
271        let f_init = || (T::zero(), T::Real::zero());
272        let f = |(acc_1, acc_2): (T, T::Real), x: T| (acc_1 + x, acc_2 + (x * x.conj()).re());
273        let f_sum = |(acc_1, acc_2): (T, T::Real), (x_1, x_2)| (acc_1 + x_1, acc_2 + x_2);
274        let f_out = |(acc_1, acc_2): (T, T::Real)| {
275            let size_1 = T::from_usize(size).unwrap();
276            let size_2 = T::Real::from_usize(size).unwrap();
277            let mean = acc_1 / size_1;
278            let var = acc_2 / size_2 - (mean * mean.conj()).re();
279            var.sqrt()
280        };
281
282        let (out, layout_out) = reduce_axes_difftype_cpu_serial(a, &la.to_dim()?, axes, f_init, f, f_sum, f_out)?;
283
284        Ok((Storage::new(out.into(), self.clone()), layout_out))
285    }
286}
287
288impl<T, D> OpL2NormAPI<T, D> for DeviceCpuSerial
289where
290    T: Clone + ComplexFloat + FromPrimitive,
291    T::Real: Clone + ComplexFloat + FromPrimitive,
292    D: DimAPI,
293{
294    type TOut = T::Real;
295
296    fn l2_norm_all(&self, a: &Vec<T>, la: &Layout<D>) -> Result<T::Real> {
297        let f_init = || T::Real::zero();
298        let f = |acc: T::Real, x: T| acc + (x * x.conj()).re();
299        let f_sum = |acc: T::Real, x: T::Real| acc + x;
300        let f_out = |acc: T::Real| acc.sqrt();
301
302        let result = reduce_all_cpu_serial(a, la, f_init, f, f_sum, f_out)?;
303        Ok(result)
304    }
305
306    fn l2_norm_axes(
307        &self,
308        a: &Vec<T>,
309        la: &Layout<D>,
310        axes: &[isize],
311    ) -> Result<(Storage<DataOwned<Vec<T::Real>>, T::Real, Self>, Layout<IxD>)> {
312        let f_init = || T::Real::zero();
313        let f = |acc: T::Real, x: T| acc + (x * x.conj()).re();
314        let f_sum = |acc: T::Real, x: T::Real| acc + x;
315        let f_out = |acc: T::Real| acc.sqrt();
316
317        let (out, layout_out) = reduce_axes_cpu_serial(a, &la.to_dim()?, axes, f_init, f, f_sum, f_out)?;
318
319        Ok((Storage::new(out.into(), self.clone()), layout_out))
320    }
321}
322
323impl<T, D> OpArgMinAPI<T, D> for DeviceCpuSerial
324where
325    T: Clone + PartialOrd,
326    D: DimAPI,
327{
328    type TOut = usize;
329
330    fn argmin_axes(
331        &self,
332        a: &Vec<T>,
333        la: &Layout<D>,
334        axes: &[isize],
335    ) -> Result<(Storage<DataOwned<Vec<usize>>, Self::TOut, Self>, Layout<IxD>)> {
336        let f_comp = |x: Option<T>, y: T| -> Option<bool> {
337            if let Some(x) = x {
338                Some(y < x)
339            } else {
340                Some(true)
341            }
342        };
343        let f_eq = |x: Option<T>, y: T| -> Option<bool> {
344            if let Some(x) = x {
345                Some(y == x)
346            } else {
347                Some(false)
348            }
349        };
350        let (out, layout_out) = reduce_axes_arg_cpu_serial(a, la, axes, f_comp, f_eq, RowMajor)?;
351        Ok((Storage::new(out.into(), self.clone()), layout_out))
352    }
353
354    fn argmin_all(&self, a: &Vec<T>, la: &Layout<D>) -> Result<Self::TOut> {
355        let f_comp = |x: Option<T>, y: T| -> Option<bool> {
356            if let Some(x) = x {
357                Some(y < x)
358            } else {
359                Some(true)
360            }
361        };
362        let f_eq = |x: Option<T>, y: T| -> Option<bool> {
363            if let Some(x) = x {
364                Some(y == x)
365            } else {
366                Some(false)
367            }
368        };
369        let result = reduce_all_arg_cpu_serial(a, la, f_comp, f_eq, RowMajor)?;
370        Ok(result)
371    }
372}
373
374impl<T, D> OpArgMaxAPI<T, D> for DeviceCpuSerial
375where
376    T: Clone + PartialOrd,
377    D: DimAPI,
378{
379    type TOut = usize;
380
381    fn argmax_axes(
382        &self,
383        a: &Vec<T>,
384        la: &Layout<D>,
385        axes: &[isize],
386    ) -> Result<(Storage<DataOwned<Vec<usize>>, Self::TOut, Self>, Layout<IxD>)> {
387        let f_comp = |x: Option<T>, y: T| -> Option<bool> {
388            if let Some(x) = x {
389                Some(y > x)
390            } else {
391                Some(true)
392            }
393        };
394        let f_eq = |x: Option<T>, y: T| -> Option<bool> {
395            if let Some(x) = x {
396                Some(y == x)
397            } else {
398                Some(false)
399            }
400        };
401        let (out, layout_out) = reduce_axes_arg_cpu_serial(a, la, axes, f_comp, f_eq, RowMajor)?;
402        Ok((Storage::new(out.into(), self.clone()), layout_out))
403    }
404
405    fn argmax_all(&self, a: &Vec<T>, la: &Layout<D>) -> Result<Self::TOut> {
406        let f_comp = |x: Option<T>, y: T| -> Option<bool> {
407            if let Some(x) = x {
408                Some(y > x)
409            } else {
410                Some(true)
411            }
412        };
413        let f_eq = |x: Option<T>, y: T| -> Option<bool> {
414            if let Some(x) = x {
415                Some(y == x)
416            } else {
417                Some(false)
418            }
419        };
420        let result = reduce_all_arg_cpu_serial(a, la, f_comp, f_eq, RowMajor)?;
421        Ok(result)
422    }
423}
424
425impl<D> OpAllAPI<bool, D> for DeviceCpuSerial
426where
427    D: DimAPI,
428{
429    type TOut = bool;
430
431    fn all_all(&self, a: &Vec<bool>, la: &Layout<D>) -> Result<bool> {
432        let f_init = || true;
433        let f = |acc, x| acc && x;
434        let f_sum = |acc1, acc2| acc1 && acc2;
435        let f_out = |acc| acc;
436
437        reduce_all_cpu_serial(a, la, f_init, f, f_sum, f_out)
438    }
439
440    fn all_axes(
441        &self,
442        a: &Vec<bool>,
443        la: &Layout<D>,
444        axes: &[isize],
445    ) -> Result<(Storage<DataOwned<Vec<bool>>, bool, Self>, Layout<IxD>)> {
446        let f_init = || true;
447        let f = |acc, x| acc && x;
448        let f_sum = |acc1, acc2| acc1 && acc2;
449        let f_out = |acc| acc;
450
451        let (out, layout_out) = reduce_axes_cpu_serial(a, &la.to_dim()?, axes, f_init, f, f_sum, f_out)?;
452        Ok((Storage::new(out.into(), self.clone()), layout_out))
453    }
454}
455
456impl<D> OpAnyAPI<bool, D> for DeviceCpuSerial
457where
458    D: DimAPI,
459{
460    type TOut = bool;
461
462    fn any_all(&self, a: &Vec<bool>, la: &Layout<D>) -> Result<bool> {
463        let f_init = || false;
464        let f = |acc, x| acc || x;
465        let f_sum = |acc1, acc2| acc1 || acc2;
466        let f_out = |acc| acc;
467
468        reduce_all_cpu_serial(a, la, f_init, f, f_sum, f_out)
469    }
470
471    fn any_axes(
472        &self,
473        a: &Vec<bool>,
474        la: &Layout<D>,
475        axes: &[isize],
476    ) -> Result<(Storage<DataOwned<Vec<bool>>, bool, Self>, Layout<IxD>)> {
477        let f_init = || false;
478        let f = |acc, x| acc || x;
479        let f_sum = |acc1, acc2| acc1 || acc2;
480        let f_out = |acc| acc;
481
482        let (out, layout_out) = reduce_axes_cpu_serial(a, &la.to_dim()?, axes, f_init, f, f_sum, f_out)?;
483        Ok((Storage::new(out.into(), self.clone()), layout_out))
484    }
485}
486
487impl<T, D> OpCountNonZeroAPI<T, D> for DeviceCpuSerial
488where
489    T: Clone + PartialEq + Zero,
490    D: DimAPI,
491{
492    type TOut = usize;
493
494    fn count_nonzero_all(&self, a: &Vec<T>, la: &Layout<D>) -> Result<usize> {
495        let f_init = || 0;
496        let f = |acc, x| if x != T::zero() { acc + 1 } else { acc };
497        let f_sum = |acc1, acc2| acc1 + acc2;
498        let f_out = |acc| acc;
499
500        reduce_all_cpu_serial(a, la, f_init, f, f_sum, f_out)
501    }
502
503    fn count_nonzero_axes(
504        &self,
505        a: &Vec<T>,
506        la: &Layout<D>,
507        axes: &[isize],
508    ) -> Result<(Storage<DataOwned<Vec<usize>>, usize, Self>, Layout<IxD>)> {
509        let f_init = || 0;
510        let f = |acc, x| if x != T::zero() { acc + 1 } else { acc };
511        let f_sum = |acc1, acc2| acc1 + acc2;
512        let f_out = |acc| acc;
513
514        let (out, layout_out) = reduce_axes_cpu_serial(a, &la.to_dim()?, axes, f_init, f, f_sum, f_out)?;
515        Ok((Storage::new(out.into(), self.clone()), layout_out))
516    }
517}
518
519impl<T, D> OpUnraveledArgMinAPI<T, D> for DeviceCpuSerial
520where
521    T: Clone + PartialOrd,
522    D: DimAPI,
523{
524    fn unraveled_argmin_axes(
525        &self,
526        a: &Vec<T>,
527        la: &Layout<D>,
528        axes: &[isize],
529    ) -> Result<(Storage<DataOwned<Vec<IxD>>, IxD, Self>, Layout<IxD>)> {
530        let f_comp = |x: Option<T>, y: T| -> Option<bool> {
531            if let Some(x) = x {
532                Some(y < x)
533            } else {
534                Some(true)
535            }
536        };
537        let f_eq = |x: Option<T>, y: T| -> Option<bool> {
538            if let Some(x) = x {
539                Some(y == x)
540            } else {
541                Some(false)
542            }
543        };
544        let (out, layout_out) = reduce_axes_unraveled_arg_cpu_serial(a, la, axes, f_comp, f_eq)?;
545        Ok((Storage::new(out.into(), self.clone()), layout_out))
546    }
547
548    fn unraveled_argmin_all(&self, a: &<Self as DeviceRawAPI<T>>::Raw, la: &Layout<D>) -> Result<D> {
549        let f_comp = |x: Option<T>, y: T| -> Option<bool> {
550            if let Some(x) = x {
551                Some(y < x)
552            } else {
553                Some(true)
554            }
555        };
556        let f_eq = |x: Option<T>, y: T| -> Option<bool> {
557            if let Some(x) = x {
558                Some(y == x)
559            } else {
560                Some(false)
561            }
562        };
563        let result = reduce_all_unraveled_arg_cpu_serial(a, la, f_comp, f_eq)?;
564        Ok(result)
565    }
566}
567
568impl<T, D> OpUnraveledArgMaxAPI<T, D> for DeviceCpuSerial
569where
570    T: Clone + PartialOrd,
571    D: DimAPI,
572{
573    fn unraveled_argmax_axes(
574        &self,
575        a: &Vec<T>,
576        la: &Layout<D>,
577        axes: &[isize],
578    ) -> Result<(Storage<DataOwned<Vec<IxD>>, IxD, Self>, Layout<IxD>)> {
579        let f_comp = |x: Option<T>, y: T| -> Option<bool> {
580            if let Some(x) = x {
581                Some(y > x)
582            } else {
583                Some(true)
584            }
585        };
586        let f_eq = |x: Option<T>, y: T| -> Option<bool> {
587            if let Some(x) = x {
588                Some(y == x)
589            } else {
590                Some(false)
591            }
592        };
593        let (out, layout_out) = reduce_axes_unraveled_arg_cpu_serial(a, la, axes, f_comp, f_eq)?;
594        Ok((Storage::new(out.into(), self.clone()), layout_out))
595    }
596
597    fn unraveled_argmax_all(&self, a: &<Self as DeviceRawAPI<T>>::Raw, la: &Layout<D>) -> Result<D> {
598        let f_comp = |x: Option<T>, y: T| -> Option<bool> {
599            if let Some(x) = x {
600                Some(y > x)
601            } else {
602                Some(true)
603            }
604        };
605        let f_eq = |x: Option<T>, y: T| -> Option<bool> {
606            if let Some(x) = x {
607                Some(y == x)
608            } else {
609                Some(false)
610            }
611        };
612        let result = reduce_all_unraveled_arg_cpu_serial(a, la, f_comp, f_eq)?;
613        Ok(result)
614    }
615}
616
617impl<D> OpSumBoolAPI<D> for DeviceCpuSerial
618where
619    D: DimAPI,
620{
621    fn sum_all(&self, a: &Vec<bool>, la: &Layout<D>) -> Result<usize> {
622        let f_init = || 0;
623        let f = |acc, x| match x {
624            true => acc + 1,
625            false => acc,
626        };
627        let f_sum = |acc1, acc2| acc1 + acc2;
628        let f_out = |acc| acc;
629
630        reduce_all_cpu_serial(a, la, f_init, f, f_sum, f_out)
631    }
632
633    fn sum_axes(
634        &self,
635        a: &Vec<bool>,
636        la: &Layout<D>,
637        axes: &[isize],
638    ) -> Result<(Storage<DataOwned<Vec<usize>>, usize, Self>, Layout<IxD>)> {
639        let f_init = || 0;
640        let f = |acc, x| match x {
641            true => acc + 1,
642            false => acc,
643        };
644        let f_sum = |acc1, acc2| acc1 + acc2;
645        let f_out = |acc| acc;
646
647        let (out, layout_out) = reduce_axes_cpu_serial(a, &la.to_dim()?, axes, f_init, f, f_sum, f_out)?;
648        Ok((Storage::new(out.into(), self.clone()), layout_out))
649    }
650}
651
652impl<TA, TB, TE, D> OpAllCloseAPI<TA, TB, TE, D> for DeviceCpuSerial
653where
654    TA: Clone + DTypePromoteAPI<TB>,
655    TB: Clone,
656    <TA as DTypePromoteAPI<TB>>::Res: ExtNum<AbsOut: DTypeCastAPI<TE>>,
657    TE: ExtFloat + Add<TE, Output = TE> + Mul<TE, Output = TE> + PartialOrd + Clone,
658    D: DimAPI,
659{
660    fn allclose_all(
661        &self,
662        a: &<Self as DeviceRawAPI<TA>>::Raw,
663        la: &Layout<D>,
664        b: &<Self as DeviceRawAPI<TB>>::Raw,
665        lb: &Layout<D>,
666        isclose_args: &IsCloseArgs<TE>,
667    ) -> Result<bool> {
668        use rstsr_dtype_traits::isclose;
669
670        if la.size() == 0 || lb.size() == 0 {
671            rstsr_raise!(InvalidValue, "zero-size array is not supported for allclose")?;
672        }
673
674        let f_init = || true;
675        let f = |acc: bool, (a_elem, b_elem): (TA, TB)| {
676            let result = isclose(&a_elem, &b_elem, isclose_args);
677            acc && result
678        };
679        let f_sum = |acc1: bool, acc2: bool| acc1 && acc2;
680        let f_out = |acc: bool| acc;
681
682        reduce_all_binary_cpu_serial(a, la, b, lb, f_init, f, f_sum, f_out)
683    }
684
685    fn allclose_axes(
686        &self,
687        _a: &<Self as DeviceRawAPI<TA>>::Raw,
688        _la: &Layout<D>,
689        _b: &<Self as DeviceRawAPI<TB>>::Raw,
690        _lb: &Layout<D>,
691        _axes: &[isize],
692        _isclose_args: &IsCloseArgs<TE>,
693    ) -> Result<(Storage<DataOwned<<Self as DeviceRawAPI<bool>>::Raw>, bool, Self>, Layout<IxD>)> {
694        unimplemented!("This function (`allclose_axes`) is not planned to be implemented yet.");
695    }
696}
rstsr_core/device_cpu_serial/reduction.rs

rstsr_core/device_cpu_serial/
reduction.rs
