tract_core/ops/nn/
reduce.rs

1use crate::internal::Axis;
2use crate::internal::*;
3use crate::ops::binary::TypedBinOp;
4use crate::ops::cast::cast;
5use crate::ops::change_axes::wire_with_rank_broadcast;
6use crate::ops::element_wise::ElementWiseOp;
7use crate::ops::math::{Mul, Square, div, square};
8use std::convert::TryFrom;
9use std::iter::Sum;
10use std::mem::transmute;
11use tract_data::internal::ClampCast;
12use tract_data::itertools::Itertools;
13use tract_linalg::routines::Func;
14use tract_ndarray::prelude::*;
15use tract_num_traits::{AsPrimitive, Bounded};
16
17macro_rules! r {
18    ($($path:ident)::* ($dt:expr) ($($args:expr),*)) => {
19        match $dt {
20            DatumType::U8   => $($path)::*::<u8,_,_,_>($($args),*),
21            DatumType::I8   => $($path)::*::<i8,_,_,_>($($args),*),
22            DatumType::U16  => $($path)::*::<u16,_,_,_>($($args),*),
23            DatumType::I16  => $($path)::*::<i16,_,_,_>($($args),*),
24            DatumType::I32  => $($path)::*::<i32,_,_,_>($($args),*),
25            DatumType::I64  => $($path)::*::<i64,_,_,_>($($args),*),
26            DatumType::F16  => $($path)::*::<f16,_,_,_>($($args),*),
27            DatumType::F32  => $($path)::*::<f32,_,_,_>($($args),*),
28            DatumType::F64  => $($path)::*::<f64,_,_,_>($($args),*),
29            DatumType::QI8(_)  => $($path)::*::<i8,_,_,_>($($args),*),
30            DatumType::QU8(_)  => $($path)::*::<u8,_,_,_>($($args),*),
31            _ => bail!("{:?} is not a number", $dt)
32        }
33    };
34    ($($path:ident)::* ($dt:expr) ($($args:expr),*); $($q_path:ident)::* ($($q_args:expr),*)) => {
35        match $dt {
36            DatumType::U8   => $($path)::*::<u8,_,_,_>($($args),*),
37            DatumType::I8   => $($path)::*::<i8,_,_,_>($($args),*),
38            DatumType::U16  => $($path)::*::<u16,_,_,_>($($args),*),
39            DatumType::I16  => $($path)::*::<i16,_,_,_>($($args),*),
40            DatumType::I32  => $($path)::*::<i32,_,_,_>($($args),*),
41            DatumType::I64  => $($path)::*::<i64,_,_,_>($($args),*),
42            DatumType::F16  => $($path)::*::<f16,_,_,_>($($args),*),
43            DatumType::F32  => $($path)::*::<f32,_,_,_>($($args),*),
44            DatumType::F64  => $($path)::*::<f64,_,_,_>($($args),*),
45            DatumType::QI8(_)  => $($q_path)::*::<i8,_,_,_>($($q_args),*),
46            DatumType::QU8(_)  => $($q_path)::*::<u8,_,_,_>($($q_args),*),
47            _ => bail!("{:?} is not a number", $dt)
48        }
49    }
50}
51
52#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
53pub enum Reducer {
54    ArgMax(bool), // take last
55    ArgMin(bool),
56    Max,
57    Min,
58    Prod,
59    Sum,
60    MeanOfSquares,
61    All,
62    Any,
63}
64
65impl Reducer {
66    pub fn reduce(&self, axes: &[usize], input: &Tensor) -> TractResult<Tensor> {
67        use Reducer::*;
68        let dt = input.datum_type();
69        let output_shape: Vec<usize> = input
70            .shape()
71            .iter()
72            .enumerate()
73            .map(|(ax, &d)| if axes.contains(&ax) { 1 } else { d })
74            .collect();
75        let (zp, scale) = input.datum_type().zp_scale();
76        unsafe {
77            let mut t = match self {
78                ArgMax(last) => {
79                    r!(Self::reduce_t(dt)(self, axes, &output_shape, input, argmax_t, *last))
80                }
81                ArgMin(last) => {
82                    r!(Self::reduce_t(dt)(self, axes, &output_shape, input, argmin_t, *last))
83                }
84                Min => r!(Self::reduce_t(dt)(self, axes, &output_shape, input, min_t, ())),
85                Max => r!(Self::reduce_t(dt)(self, axes, &output_shape, input, max_t, ())),
86                Prod => {
87                    r!(Self::reduce_t(dt)(self, axes, &output_shape, input, prod_t, ()); Self::reduce_t(self, axes, &output_shape, input, q_prod_t, (zp, scale)))
88                }
89                Sum => {
90                    if dt.is_float() {
91                        dispatch_floatlike!(Self::sum(dt)(self, axes, input))
92                    } else {
93                        r!(Self::reduce_t(dt)(
94                            self,
95                            axes,
96                            &output_shape,
97                            input,
98                            q_sum_t,
99                            (zp, scale)
100                        ))
101                    }
102                }
103                MeanOfSquares => self.mean_of_squares(axes, input)?,
104                All => Self::reduce_t(self, axes, &output_shape, input, all_bool, ()),
105                Any => Self::reduce_t(self, axes, &output_shape, input, any_bool, ()),
106            };
107            if input.datum_type().is_quantized()
108                && input.datum_type().unquantized() == t.datum_type().unquantized()
109            {
110                t.set_datum_type(input.datum_type());
111            }
112            Ok(t)
113        }
114    }
115
116    unsafe fn reduce_t<T, TO, F, A>(
117        &self,
118        axes: &[usize],
119        output_shape: &[usize],
120        input_tensor: &Tensor,
121        f: F,
122        args: A,
123    ) -> Tensor
124    where
125        F: for<'a> Fn(ArrayViewD<'a, T>, A) -> TO,
126        T: Copy + Datum,
127        TO: Copy + Datum,
128        A: Copy,
129    {
130        use ndarray::*;
131        let input = unsafe { input_tensor.to_array_view_unchecked::<T>() };
132        let result = Array::from_shape_fn(output_shape, |coords| {
133            let slice_spec: Vec<SliceInfoElem> = coords
134                .slice()
135                .iter()
136                .enumerate()
137                .map(|(ax, &d)| if axes.contains(&ax) { (..).into() } else { d.into() })
138                .collect();
139            let slice_info = SliceInfo::<_, IxDyn, IxDyn>::try_from(slice_spec).unwrap();
140            let slice = input.slice(&slice_info);
141            f(slice, args)
142        });
143        result.into_tensor()
144    }
145
146    // sum is a special citizen: enough activity that it gets "special"
147    // treatment. we could use the same "algo" for min, max and prod, to the
148    // price of more code in the library. argmax and argmin are more
149    // tricky (not associative)
150    unsafe fn sum<T>(&self, axes: &[usize], input: &Tensor) -> Tensor
151    where
152        T: Copy + Datum + num_traits::Zero + Sum,
153        f16: AsPrimitive<T>,
154        f32: AsPrimitive<T>,
155    {
156        if axes.len() == 0 {
157            return input.to_owned();
158        }
159
160        // use tract-optimized path only when single reuction axis and is at end
161        if axes.len() > 1 || axes[0] != input.rank() - 1 {
162            let mut operative_axes = vec![];
163            let mut operative_shape: Vec<usize> = vec![];
164            for (ix, dim) in input.shape().iter().enumerate() {
165                // axis is reduced, but is not the first of a series of reduced axes
166                if ix > 0 && axes.contains(&ix) && axes.contains(&(ix - 1)) {
167                    *operative_shape.last_mut().unwrap() *= *dim;
168                } else if axes.contains(&ix) {
169                    operative_axes.push(operative_shape.len());
170                    operative_shape.push(*dim);
171                } else {
172                    operative_shape.push(*dim);
173                }
174            }
175            let mut output = unsafe {
176                input
177                    .to_array_view_unchecked::<T>()
178                    .into_shape_with_order(operative_shape)
179                    .unwrap()
180                    .sum_axis(Axis(*operative_axes.iter().max().unwrap()))
181            };
182
183            for axis in operative_axes.iter().rev().skip(1) {
184                output = output.sum_axis(Axis(*axis));
185            }
186
187            let mut output = output.into_tensor();
188
189            for &axis in axes {
190                output.insert_axis(axis).unwrap();
191            }
192
193            output
194        } else {
195            let mut output: Option<ArrayD<T>> = None;
196            for axis in axes.iter().copied() {
197                let input_view = output
198                    .as_ref()
199                    .map(|o| o.view())
200                    .unwrap_or_else(|| unsafe { input.to_array_view_unchecked::<T>() });
201
202                // Create array that will contain intermidiate result
203                let reduced_dim = input_view.shape()[axis];
204                let input_stride = input_view.strides()[axis] as usize;
205                let output_shape = input_view
206                    .shape()
207                    .iter()
208                    .enumerate()
209                    .map(|(idx, dim)| if idx != axis { *dim } else { 1 })
210                    .collect_vec();
211
212                output = Some(if let Some(full) = input_view.as_slice() {
213                    // Whole input is C-contiguous and `axis` is the last axis,
214                    // so it lays out as [n_rows, reduced_dim] row-major: sum
215                    // each row in one pass. Rows split across threads while
216                    // each row's reduction stays serial and bit-identical.
217                    let n_rows = full.len() / reduced_dim;
218                    let mut out = vec![T::zero(); n_rows];
219                    let total = full.len();
220                    // Reduce kernels are Send + Sync; build once and share by ref.
221                    let sum_f16 = Func::ReduceSum.reduce_f16().unwrap();
222                    let sum_f32 = Func::ReduceSum.reduce_f32().unwrap();
223                    tract_linalg::multithread::par_chunks_mut(
224                        &mut out,
225                        1,
226                        total,
227                        |first_row, o| {
228                            let rows = full[first_row * reduced_dim..][..o.len() * reduced_dim]
229                                .chunks_exact(reduced_dim);
230                            if reduced_dim >= 4 && T::datum_type() == f16::datum_type() {
231                                for (x, c) in o.iter_mut().zip(rows) {
232                                    let c: &[f16] = unsafe { std::mem::transmute(c) };
233                                    *x = sum_f16.run_with_params(c, ())?.as_();
234                                }
235                            } else if reduced_dim >= 4 && T::datum_type() == f32::datum_type() {
236                                for (x, c) in o.iter_mut().zip(rows) {
237                                    let c: &[f32] = unsafe { std::mem::transmute(c) };
238                                    *x = sum_f32.run_with_params(c, ())?.as_();
239                                }
240                            } else {
241                                // reduced_dim < 4 (kernel dispatch not worth it) or a
242                                // non-f16/f32 type: a plain sum matches the kernel's
243                                // remainder path bit-for-bit.
244                                for (x, c) in o.iter_mut().zip(rows) {
245                                    *x = c.iter().cloned().sum::<T>();
246                                }
247                            }
248                            Ok(())
249                        },
250                    )
251                    .unwrap();
252                    ArrayD::from_shape_vec(output_shape.clone(), out).unwrap()
253                } else {
254                    ArrayD::from_shape_fn(output_shape.clone(), |coords| {
255                        let first: *const T = &input_view[coords];
256                        let mut sum = T::zero();
257                        for i in 0..reduced_dim {
258                            sum = sum + unsafe { *(first.add(i * input_stride)) };
259                        }
260                        sum
261                    })
262                });
263            }
264            output.unwrap().into_tensor()
265        }
266    }
267
268    fn mean_of_squares(&self, axis: &[usize], input: &Tensor) -> TractResult<Tensor> {
269        let dt = input.datum_type();
270        let mut input = input.cast_to::<f32>()?.into_owned();
271        input.try_as_plain_mut()?.as_slice_mut::<f32>()?.iter_mut().for_each(|x| *x = *x * *x);
272        let mut output = unsafe { self.sum::<f32>(axis, &input) };
273        let norm = output.len() as f32 / input.len() as f32;
274        output.try_as_plain_mut()?.as_slice_mut::<f32>()?.iter_mut().for_each(|x| *x *= norm);
275        Ok(output.cast_to_dt(dt)?.into_owned())
276    }
277}
278
279fn argmax_t<T>(v: ArrayViewD<T>, last: bool) -> i64
280where
281    T: Copy + Datum + num_traits::Bounded + ::std::cmp::PartialOrd,
282{
283    v.iter()
284        .copied()
285        .enumerate()
286        .fold(
287            (0usize, T::min_value()),
288            |acc, v| {
289                if v.1 > acc.1 || (last && acc.1 == v.1) { v } else { acc }
290            },
291        )
292        .0 as i64
293}
294
295fn argmin_t<T>(v: ArrayViewD<T>, last: bool) -> i64
296where
297    T: Copy + Datum + num_traits::Bounded + ::std::cmp::PartialOrd,
298{
299    v.iter()
300        .copied()
301        .enumerate()
302        .fold(
303            (0usize, T::max_value()),
304            |acc, v| {
305                if v.1 < acc.1 || (last && acc.1 == v.1) { v } else { acc }
306            },
307        )
308        .0 as i64
309}
310
311fn max_t<T>(v: ArrayViewD<T>, _: ()) -> T
312where
313    T: Copy + Datum + num_traits::Bounded + ::std::cmp::PartialOrd,
314{
315    if T::datum_type() == f32::datum_type()
316        && let Some(slice) = v.as_slice()
317        && !slice.is_empty()
318    {
319        let slice = unsafe { transmute::<&[T], &[f32]>(slice) };
320        let max = Func::ReduceMax.reduce_f32().unwrap().run(slice).unwrap();
321        // SAFETY: T is f32 in this branch (checked above).
322        return unsafe { std::mem::transmute_copy::<f32, T>(&max) };
323    }
324    v.fold(T::min_value(), |acc, &v| if acc > v { acc } else { v })
325}
326
327fn min_t<T>(v: ArrayViewD<T>, _: ()) -> T
328where
329    T: Copy + Datum + num_traits::Bounded + ::std::cmp::PartialOrd,
330{
331    if T::datum_type() == f32::datum_type()
332        && let Some(slice) = v.as_slice()
333        && !slice.is_empty()
334    {
335        let slice = unsafe { transmute::<&[T], &[f32]>(slice) };
336        let min = Func::ReduceMin.reduce_f32().unwrap().run(slice).unwrap();
337        // SAFETY: T is f32 in this branch (checked above).
338        return unsafe { std::mem::transmute_copy::<f32, T>(&min) };
339    }
340    v.fold(T::max_value(), |acc, &v| if acc < v { acc } else { v })
341}
342
343fn prod_t<T>(v: ArrayViewD<T>, _: ()) -> T
344where
345    T: Copy + Datum + num_traits::One,
346{
347    v.fold(T::one(), |acc, &v| acc * v)
348}
349
350fn q_prod_t<T>(v: ArrayViewD<T>, zp_scale: (i32, f32)) -> T
351where
352    T: Copy + num_traits::AsPrimitive<f32> + Bounded + Datum,
353    f32: num_traits::AsPrimitive<T>,
354{
355    let (zp, scale) = zp_scale;
356    (v.fold(1f32, |acc, &v| acc * (v.as_() - zp as f32)) * scale.powi(v.len() as i32 - 1)
357        + zp as f32)
358        .clamp_cast()
359}
360
361fn q_sum_t<T>(v: ArrayViewD<T>, zp_scale: (i32, f32)) -> T
362where
363    T: Copy + Bounded + num_traits::AsPrimitive<i32> + Datum,
364    i32: num_traits::AsPrimitive<T>,
365{
366    let (zp, _) = zp_scale;
367    (v.fold(0i32, |acc, &v| acc + v.as_()) - zp * (v.len() as i32 - 1)).clamp_cast()
368}
369
370fn all_bool(v: ArrayViewD<bool>, _: ()) -> bool {
371    v.iter().all(|v| *v)
372}
373
374fn any_bool(v: ArrayViewD<bool>, _: ()) -> bool {
375    v.iter().any(|v| *v)
376}
377
378#[derive(Clone, Debug, new, Hash, PartialEq, Eq)]
379pub struct Reduce {
380    pub axes: TVec<usize>,
381    pub reducer: Reducer,
382}
383
384impl Op for Reduce {
385    fn name(&self) -> StaticName {
386        format!("Reduce<{:?}>", self.reducer).into()
387    }
388    fn info(&self) -> TractResult<Vec<String>> {
389        Ok(vec![format!("axes: {:?}", self.axes)])
390    }
391    op_as_typed_op!();
392}
393
394impl EvalOp for Reduce {
395    op_out_of_plan!();
396
397    fn eval(&self, _ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
398        Ok(tvec!(self.reducer.reduce(&self.axes, &inputs[0])?.into()))
399    }
400}
401
402impl TypedOp for Reduce {
403    fn input_roi(
404        &self,
405        model: &TypedModel,
406        node: &TypedNode,
407    ) -> TractResult<Option<TVec<Option<TDim>>>> {
408        crate::optim::propagate_roi::bubble_roi(model, node)
409    }
410
411    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
412        ensure!(self.axes.iter().tuple_windows().all(|(a, b)| a < b));
413        if inputs[0].datum_type == TDim::datum_type() {
414            bail!("Reduce input must be cast from TDim to i64 beforehand")
415        }
416        let mut shape: TVec<_> = inputs[0].shape.to_tvec();
417        for &ax in &self.axes {
418            shape[ax] = 1.to_dim();
419        }
420        let dt = if let Reducer::ArgMax(_) | Reducer::ArgMin(_) = self.reducer {
421            DatumType::I64
422        } else {
423            inputs[0].datum_type
424        };
425        Ok(tvec!(dt.fact(shape)))
426    }
427
428    fn declutter(
429        &self,
430        model: &TypedModel,
431        node: &TypedNode,
432    ) -> TractResult<Option<TypedModelPatch>> {
433        if let Some(patch) = self.declutter_mean_of_square(model, node)? {
434            return Ok(Some(patch));
435        }
436        if let Some(patch) = self.declutter_scalar_mul_then_sum(model, node)? {
437            return Ok(Some(patch));
438        }
439        if let Some(patch) = self.declutter_reduce_reduce(model, node)? {
440            return Ok(Some(patch));
441        }
442        if let Some(patch) = super::rms_norm::detect_rms_norm(self, model, node)? {
443            return Ok(Some(patch));
444        }
445        Ok(None)
446    }
447
448    fn cost(&self, inputs: &[&TypedFact]) -> TractResult<TVec<(Cost, TDim)>> {
449        let dt = inputs[0].datum_type;
450        let count: TDim = inputs[0].shape.iter().product();
451        match self.reducer {
452            Reducer::Sum
453            | Reducer::Prod
454            | Reducer::Min
455            | Reducer::Max
456            | Reducer::All
457            | Reducer::Any => Ok(tvec!((Cost::FMA(dt), count))),
458            Reducer::MeanOfSquares => Ok(tvec!((Cost::FMA(dt), count * 2))),
459            Reducer::ArgMax(_) | Reducer::ArgMin(_) => Ok(tvec!((Cost::FMA(dt), count))),
460        }
461    }
462
463    fn axes_mapping(
464        &self,
465        inputs: &[&TypedFact],
466        outputs: &[&TypedFact],
467    ) -> TractResult<AxesMapping> {
468        let mut letters = 'a'..;
469        let axes = (0..inputs[0].rank())
470            .flat_map(|ix| {
471                if self.axes.contains(&ix) {
472                    tvec!(
473                        Axis::new(letters.next().unwrap(), inputs.len(), outputs.len())
474                            .input(0, ix),
475                        Axis::new(letters.next().unwrap(), inputs.len(), outputs.len())
476                            .output(0, ix),
477                    )
478                } else {
479                    tvec!(
480                        Axis::new(letters.next().unwrap(), inputs.len(), outputs.len())
481                            .input(0, ix)
482                            .output(0, ix)
483                    )
484                }
485                .into_iter()
486            })
487            .collect_vec();
488        AxesMapping::new(1, 1, axes)
489    }
490
491    fn change_axes(
492        &self,
493        model: &TypedModel,
494        node: &TypedNode,
495        _io: InOut,
496        change: &AxisOp,
497    ) -> TractResult<Option<AxisChangeConsequence>> {
498        let mut axes = tvec!();
499        for reduced in &self.axes {
500            rule_if_some!(axis = change.transform_axis(*reduced));
501            axes.push(axis);
502        }
503        axes.sort();
504        let op = Some(Box::new(Self { axes, ..self.clone() }) as _);
505        Ok(Some(AxisChangeConsequence::new(model, node, op, change)))
506    }
507
508    fn slice(
509        &self,
510        patch: &mut TypedModelPatch,
511        _model: &TypedModel,
512        node: &TypedNode,
513        _prefix: &str,
514        inputs: &[OutletId],
515        output_axis: usize,
516        _start: &TDim,
517        _end: &TDim,
518    ) -> TractResult<Option<TVec<OutletId>>> {
519        rule_if!(!self.axes.contains(&output_axis));
520        patch.wire_node(&node.name, &node.op, inputs).map(Some)
521    }
522
523    as_op!();
524}
525
526impl Reduce {
527    fn declutter_reduce_reduce(
528        &self,
529        model: &TypedModel,
530        node: &TypedNode,
531    ) -> TractResult<Option<TypedModelPatch>> {
532        use Reducer::*;
533        rule_if_some!(prec = model.linear_prec(node.id)?);
534        rule_if_some!(prec_reduce = prec.op_as::<Self>());
535        rule_if!(prec_reduce.reducer == self.reducer);
536        rule_if!([Sum, Prod, Min, Max].contains(&self.reducer));
537        let mut patch = TypedModelPatch::default();
538        let wire = patch.tap_model(model, prec.inputs[0])?;
539        let wire = patch.wire_node(
540            &node.name,
541            Self {
542                reducer: self.reducer,
543                axes: prec_reduce
544                    .axes
545                    .iter()
546                    .chain(self.axes.iter())
547                    .copied()
548                    .sorted()
549                    .dedup()
550                    .collect(),
551            },
552            &[wire],
553        )?;
554        patch.shunt_outside(model, node.id.into(), wire[0])?;
555        Ok(Some(patch))
556    }
557
558    fn declutter_scalar_mul_then_sum(
559        &self,
560        model: &TypedModel,
561        node: &TypedNode,
562    ) -> TractResult<Option<TypedModelPatch>> {
563        if self.reducer == Reducer::Sum {
564            rule_if_some!(prec = model.linear_prec(node.id)?);
565            rule_if_some!(prec_bin = prec.op_as::<TypedBinOp>());
566            rule_if!(prec_bin.0.is::<Mul>());
567            let mul_input_fact = model.node_input_facts(prec.id)?;
568            rule_if_some!(
569                scalar_slot = mul_input_fact
570                    .iter()
571                    .position(|f| f.konst.as_ref().is_some_and(|k| k.volume() == 1))
572            );
573            let mut patch = TypedModelPatch::default();
574            let scalar = patch.tap_model(model, prec.inputs[scalar_slot])?;
575            let wire = patch.tap_model(model, prec.inputs[1 - scalar_slot])?;
576            let wire = patch.wire_node(&node.name, self.clone(), &[wire])?[0];
577            let wire = patch.wire_node(&prec.name, prec_bin.clone(), &[wire, scalar])?[0];
578            patch.shunt_outside(model, node.id.into(), wire)?;
579            return Ok(Some(patch));
580        }
581        Ok(None)
582    }
583
584    fn declutter_mean_of_square(
585        &self,
586        model: &TypedModel,
587        node: &TypedNode,
588    ) -> TractResult<Option<TypedModelPatch>> {
589        if self.reducer == Reducer::Sum {
590            rule_if_some!(prec = model.linear_prec(node.id)?);
591            rule_if_some!(prec_ew = prec.op_as::<ElementWiseOp>());
592            rule_if!(prec_ew.0.is::<Square>());
593            rule_if!(node.outputs.len() == 1);
594            rule_if!(node.outputs[0].successors.len() == 1);
595            let our_inlet = node.outputs[0].successors[0];
596            let succ = model.node(our_inlet.node);
597            rule_if_some!(succ_bin = succ.op_as::<TypedBinOp>());
598            rule_if!(succ_bin.0.is::<Mul>());
599            let other = succ.inputs[1 - our_inlet.slot];
600            rule_if_some!(other_konst = model.outlet_fact(other)?.uniform.as_ref());
601            let norm: TDim = self.axes.iter().map(|&ax| &prec.outputs[0].fact.shape[ax]).product();
602            rule_if_some!(norm = norm.as_i64());
603            rule_if!(norm > 0);
604            let norm = tensor0((norm as f32).recip());
605            if other_konst.close_enough(&norm, Approximation::Close).is_ok() {
606                let mut patch = TypedModelPatch::default();
607                let wire = patch.tap_model(model, prec.inputs[0])?;
608                let wire = patch.wire_node(
609                    &node.name,
610                    Reduce::new(self.axes.clone(), Reducer::MeanOfSquares),
611                    &[wire],
612                )?[0];
613                patch.shunt_outside(model, succ.id.into(), wire)?;
614                return Ok(Some(patch));
615            }
616        }
617        Ok(None)
618    }
619}
620
621pub fn expand_mean_of_squares(
622    _ctx: &(),
623    model: &TypedModel,
624    node: &TypedNode,
625    name: &str,
626    op: &Reduce,
627) -> TractResult<Option<TypedModelPatch>> {
628    rule_if!(op.reducer == Reducer::MeanOfSquares);
629    let mut patch = TypedModelPatch::default();
630    let mut wire = tvec!(patch.tap_model(model, node.inputs[0])?);
631    let input_fact = model.outlet_fact(node.inputs[0])?;
632    let dt = input_fact.datum_type;
633    if dt != f32::datum_type() {
634        wire = patch.wire_node(format!("{name}.to_f32"), cast(f32::datum_type()), &wire)?;
635    }
636    wire = patch.wire_node(format!("{name}.sqr"), square(), &wire)?;
637    wire = patch.wire_node(
638        format!("{name}.sum"),
639        Reduce::new(op.axes.clone(), Reducer::Sum),
640        &wire,
641    )?;
642    let card = input_fact
643        .shape
644        .iter()
645        .enumerate()
646        .filter(|(ix, _dim)| op.axes.contains(ix))
647        .map(|(_ix, dim)| dim)
648        .product::<TDim>();
649    let card = patch.add_const(format!("{name}.card"), tensor0(card))?;
650    let card = patch.wire_node(format!("{name}.card_to_f32"), cast(f32::datum_type()), &[card])?;
651
652    wire =
653        wire_with_rank_broadcast(format!("{name}.norm"), &mut patch, div(), &[wire[0], card[0]])?;
654    if dt != f32::datum_type() {
655        wire = patch.wire_node(format!("{name}.from_f32"), cast(dt), &wire)?;
656    }
657    patch.shunt_outside(model, node.id.into(), wire[0])?;
658    Ok(Some(patch))
659}
660
661#[cfg(test)]
662mod tests {
663    use super::*;
664
665    // Guards the f32 max reduction (max_t): the SIMD `max_f32` kernel result must
666    // be returned (contiguous path), with the scalar fold used only for strided
667    // slices. Checked against explicit per-row / per-col references.
668    #[test]
669    fn reduce_max_f32_contiguous_and_strided() {
670        let (r, c) = (5usize, 37usize); // c not a multiple of the SIMD width (tail)
671        let data: Vec<f32> = (0..r * c).map(|i| ((i * 31 % 97) as f32) - 48.0).collect();
672        let t = Tensor::from_shape(&[r, c], &data).unwrap();
673
674        // axis 1: per-row max — contiguous slices -> SIMD path.
675        let got = Reducer::Max.reduce(&[1], &t).unwrap();
676        assert_eq!(got.shape(), &[r, 1]);
677        for (i, &g) in unsafe { got.as_slice_unchecked::<f32>() }.iter().enumerate() {
678            let want = data[i * c..(i + 1) * c].iter().copied().fold(f32::MIN, f32::max);
679            assert_eq!(g, want, "row {i}");
680        }
681
682        // axis 0: per-col max — strided slices -> scalar fold.
683        let got = Reducer::Max.reduce(&[0], &t).unwrap();
684        assert_eq!(got.shape(), &[1, c]);
685        for (j, &g) in unsafe { got.as_slice_unchecked::<f32>() }.iter().enumerate() {
686            let want = (0..r).map(|i| data[i * c + j]).fold(f32::MIN, f32::max);
687            assert_eq!(g, want, "col {j}");
688        }
689
690        // k == 1 (single-element reduction) exercises the SIMD-path length guard.
691        let t1 = Tensor::from_shape(&[3, 1], &[1.0f32, -2.0, 3.0]).unwrap();
692        let got = Reducer::Max.reduce(&[1], &t1).unwrap();
693        assert_eq!(unsafe { got.as_slice_unchecked::<f32>() }, &[1.0, -2.0, 3.0]);
694    }
695
696    // Same coverage for the f32 min reduction (min_t -> SIMD min_f32 / scalar fold).
697    #[test]
698    fn reduce_min_f32_contiguous_and_strided() {
699        let (r, c) = (5usize, 37usize); // c not a multiple of the SIMD width (tail)
700        let data: Vec<f32> = (0..r * c).map(|i| ((i * 31 % 97) as f32) - 48.0).collect();
701        let t = Tensor::from_shape(&[r, c], &data).unwrap();
702
703        // axis 1: per-row min — contiguous slices -> SIMD path.
704        let got = Reducer::Min.reduce(&[1], &t).unwrap();
705        assert_eq!(got.shape(), &[r, 1]);
706        for (i, &g) in unsafe { got.as_slice_unchecked::<f32>() }.iter().enumerate() {
707            let want = data[i * c..(i + 1) * c].iter().copied().fold(f32::MAX, f32::min);
708            assert_eq!(g, want, "row {i}");
709        }
710
711        // axis 0: per-col min — strided slices -> scalar fold.
712        let got = Reducer::Min.reduce(&[0], &t).unwrap();
713        assert_eq!(got.shape(), &[1, c]);
714        for (j, &g) in unsafe { got.as_slice_unchecked::<f32>() }.iter().enumerate() {
715            let want = (0..r).map(|i| data[i * c + j]).fold(f32::MAX, f32::min);
716            assert_eq!(g, want, "col {j}");
717        }
718    }
719}
tract_core/ops/nn/reduce.rs

tract_core/ops/nn/
reduce.rs
