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

tract_core/ops/nn/
reduce.rs
