Skip to main content

tract_core/ops/cnn/
sumpool.rs

1use crate::internal::*;
2use num_traits::AsPrimitive;
3use std::iter::Sum;
4
5use crate::ops::cnn::pools::{ConcretePoolGeometry, PoolGeometry, PoolSpec};
6
7crate::declare_knob!(
8    TRACT_AVGPOOL_SEPARABLE,
9    bool,
10    false,
11    "Use the separable average-pool kernel for stride-1 NCHW/NHWC pools. Not bit-identical: \
12     it reassociates the sum, permitted by SumPool's Validation::Rounding contract."
13);
14
15#[derive(Debug, Clone, new, Hash, PartialEq, Eq)]
16pub struct SumPool {
17    pub pool_spec: PoolSpec,
18    pub count_include_pad: bool,
19    pub normalize: bool,
20}
21
22impl Op for SumPool {
23    fn name(&self) -> StaticName {
24        "SumPool".into()
25    }
26
27    fn info(&self) -> TractResult<Vec<String>> {
28        Ok(self.pool_spec.info())
29    }
30
31    fn validation(&self) -> Validation {
32        Validation::Rounding
33    }
34
35    op_as_typed_op!();
36}
37
38impl EvalOp for SumPool {
39    fn is_stateless(&self) -> bool {
40        true
41    }
42
43    fn eval(&self, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
44        let shape: TVec<TDim> = inputs[0].shape().iter().map(|d| d.to_dim()).collect();
45        self.to_optimized(&shape)?.eval(inputs)
46    }
47}
48
49impl TypedOp for SumPool {
50    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
51        self.pool_spec.output_facts(inputs)
52    }
53
54    fn declutter(
55        &self,
56        model: &TypedModel,
57        node: &TypedNode,
58    ) -> TractResult<Option<TypedModelPatch>> {
59        let fact = model.outlet_fact(node.inputs[0])?;
60        if let Some(pool_spec) = self.pool_spec.declutter(&fact.shape)? {
61            return Ok(Some(TypedModelPatch::replace_single_op(
62                model,
63                node,
64                &node.inputs,
65                Self { pool_spec, ..self.clone() },
66            )?));
67        }
68        Ok(None)
69    }
70
71    /// Lower to `OptSumPool` with the geometry pre-resolved to `Concrete` when the
72    /// input shape is fixed, so the `Patch` is built once here rather than per eval.
73    /// Symbolic shapes are left as `SumPool`.
74    fn codegen(
75        &self,
76        model: &TypedModel,
77        node: &TypedNode,
78    ) -> TractResult<Option<TypedModelPatch>> {
79        let fact = model.outlet_fact(node.inputs[0])?;
80        if fact.shape.as_concrete().is_none() {
81            return Ok(None);
82        }
83        let mut op = self.to_optimized(&fact.shape.to_tvec())?;
84        op.geometry = op.geometry.optimize_if(fact.shape.as_concrete())?;
85        Ok(Some(TypedModelPatch::replace_single_op(model, node, &node.inputs, op)?))
86    }
87
88    as_op!();
89}
90
91impl SumPool {
92    fn to_optimized(&self, input_shape: &[TDim]) -> TractResult<OptSumPool> {
93        Ok(OptSumPool {
94            pool_spec: self.pool_spec.clone(),
95            count_include_pad: self.count_include_pad,
96            normalize: self.normalize,
97            geometry: self.pool_spec.compute_geo(input_shape)?,
98        })
99    }
100}
101
102#[derive(Debug, Clone, new, Hash, PartialEq, Eq)]
103pub struct OptSumPool {
104    pub pool_spec: PoolSpec,
105    pub count_include_pad: bool,
106    pub normalize: bool,
107    pub geometry: PoolGeometry,
108}
109
110impl Op for OptSumPool {
111    fn name(&self) -> StaticName {
112        "OptSumPool".into()
113    }
114
115    fn info(&self) -> TractResult<Vec<String>> {
116        Ok(self.pool_spec.info())
117    }
118
119    fn validation(&self) -> Validation {
120        Validation::Rounding
121    }
122
123    op_as_typed_op!();
124}
125
126impl EvalOp for OptSumPool {
127    fn is_stateless(&self) -> bool {
128        true
129    }
130
131    fn eval(&self, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
132        let input = args_1!(inputs);
133        let geo = self.geometry.to_concrete(input.shape())?;
134        let values = if input.datum_type().is_float() {
135            let mut values =
136                unsafe { Tensor::uninitialized_dt(input.datum_type(), &geo.output_shape.shape)? };
137            dispatch_floatlike!(Self::eval_t(input.datum_type())(
138                self,
139                &*input,
140                values.as_ptr_mut()?,
141                geo.as_ref()
142            ))?;
143            values
144        } else {
145            let mut values =
146                unsafe { Tensor::uninitialized_dt(DatumType::F32, &geo.output_shape.shape)? };
147            let input_f32 = input.cast_to_dt(DatumType::F32)?;
148            self.eval_t::<f32>(input_f32.as_ref(), values.as_ptr_mut()?, geo.as_ref())?;
149            values.cast_to_dt(input.datum_type())?.into_owned()
150        };
151
152        Ok(tvec!(values.into_tvalue()))
153    }
154}
155
156impl TypedOp for OptSumPool {
157    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
158        self.pool_spec.output_facts(inputs)
159    }
160
161    fn declutter(
162        &self,
163        model: &TypedModel,
164        node: &TypedNode,
165    ) -> TractResult<Option<TypedModelPatch>> {
166        let fact = model.outlet_fact(node.inputs[0])?;
167        if let Some(pool_spec) = self.pool_spec.declutter(&fact.shape)? {
168            return Ok(Some(TypedModelPatch::replace_single_op(
169                model,
170                node,
171                &node.inputs,
172                Self { pool_spec, ..self.clone() },
173            )?));
174        }
175        Ok(None)
176    }
177
178    as_op!();
179}
180
181impl OptSumPool {
182    fn eval_t<T: Copy + Datum + Sum + num_traits::Float>(
183        &self,
184        input: &Tensor,
185        values_ptr: *mut T,
186        geo: &ConcretePoolGeometry,
187    ) -> TractResult<()>
188    where
189        usize: AsPrimitive<T>,
190    {
191        if self.try_fast_2d::<T>(input, values_ptr, geo)? {
192            return Ok(());
193        }
194        let input_ptr = input.as_ptr::<T>()?;
195
196        let n = *geo.input_shape.n().unwrap_or(&1);
197        let n_stride_i = geo.input_shape.n_stride().unwrap_or(&0);
198        let n_stride_o = geo.output_shape.n_stride().unwrap_or(&0);
199        unsafe {
200            geo.patch.visit_output(|visitor| {
201                let div: Option<T> = if self.normalize {
202                    Some(
203                        if self.count_include_pad {
204                            geo.patch.standard_layout_data_field.len().as_()
205                        } else {
206                            visitor.valid_count().as_()
207                        }
208                        .recip(),
209                    )
210                } else {
211                    None
212                };
213                for n in 0..n {
214                    let input_offset = n * n_stride_i;
215                    let output_offset = n * n_stride_o;
216                    for c in 0..*geo.input_shape.c() {
217                        let input_offset = input_offset + geo.input_shape.c_stride() * c;
218                        let output_offset = output_offset + geo.output_shape.c_stride() * c;
219                        let sum = visitor
220                            .valid_offsets()
221                            .map(|v| *input_ptr.offset(v + input_offset as isize))
222                            .sum::<T>();
223
224                        if let Some(div) = div {
225                            *values_ptr.offset(output_offset as isize + visitor.output_offset) =
226                                sum * div;
227                        }
228                    }
229                }
230            });
231        }
232        Ok(())
233    }
234
235    /// Opt-in separable average-pool fast path, gated by `TRACT_AVGPOOL_SEPARABLE`.
236    /// Returns `true` when it handled the eval, `false` to fall back to the generic
237    /// kernel. Restricted to rank-2, stride-1, dilation-1, normalize pools; the NCHW
238    /// and NHWC layouts have dedicated kernels, any other layout falls back.
239    fn try_fast_2d<T: Copy + Datum + num_traits::Float>(
240        &self,
241        input: &Tensor,
242        values_ptr: *mut T,
243        geo: &ConcretePoolGeometry,
244    ) -> TractResult<bool>
245    where
246        usize: AsPrimitive<T>,
247    {
248        let patch = &geo.patch;
249        if !TRACT_AVGPOOL_SEPARABLE.get()
250            || !self.normalize
251            || patch.rank() != 2
252            || *patch.spec.strides != [1, 1]
253            || *patch.spec.dilations != [1, 1]
254        {
255            return Ok(false);
256        }
257        let input_ptr = input.as_ptr::<T>()?;
258        let ish = &geo.input_shape;
259        if *ish.w_stride() == 1 {
260            unsafe {
261                self.fast_2d_separable::<T>(input_ptr, values_ptr, geo);
262            }
263            Ok(true)
264        } else if *ish.c_stride() == 1 && *ish.w_stride() == *ish.c() {
265            unsafe {
266                self.fast_2d_separable_nhwc::<T>(input_ptr, values_ptr, geo);
267            }
268            Ok(true)
269        } else {
270            Ok(false)
271        }
272    }
273
274    /// Separable running-sum average pool. Out-of-bounds taps contribute 0, so the
275    /// box sum over the padded window equals the sum of in-bounds values; the divisor
276    /// is the per-cell valid count, itself separable into `kx_valid * ky_valid`.
277    /// Reassociates the sum, so it is not bit-identical to the generic kernel.
278    unsafe fn fast_2d_separable<T: Copy + Datum + num_traits::Float>(
279        &self,
280        input_ptr: *const T,
281        values_ptr: *mut T,
282        geo: &ConcretePoolGeometry,
283    ) where
284        usize: AsPrimitive<T>,
285    {
286        let ish = &geo.input_shape;
287        let osh = &geo.output_shape;
288        let (h, w) = (ish.hw_dims()[0] as isize, ish.hw_dims()[1] as isize);
289        let (ho, wo) = (geo.patch.output_shape[0], geo.patch.output_shape[1]);
290        let (kh, kw) =
291            (geo.patch.spec.kernel_shape[0] as isize, geo.patch.spec.kernel_shape[1] as isize);
292        let (pt, pl) = (geo.patch.pad_before[0] as isize, geo.patch.pad_before[1] as isize);
293        let ih_stride = *ish.h_stride() as isize;
294        let oh_stride = *osh.h_stride() as isize;
295        let ow_stride = *osh.w_stride() as isize;
296        let n = *ish.n().unwrap_or(&1);
297        let in_stride = *ish.n_stride().unwrap_or(&0) as isize;
298        let on_stride = *osh.n_stride().unwrap_or(&0) as isize;
299        let c = *ish.c();
300        let ic_stride = *ish.c_stride() as isize;
301        let oc_stride = *osh.c_stride() as isize;
302
303        let axis_valid = |out: usize, k: isize, pad: isize, lim: isize| -> Vec<usize> {
304            (0..out)
305                .map(|o| {
306                    let lo = o as isize - pad;
307                    let start = (-lo).max(0);
308                    let end = (lim - lo).min(k);
309                    (end - start).max(0) as usize
310                })
311                .collect()
312        };
313        let kx_valid = axis_valid(wo, kw, pl, w);
314        let ky_valid = axis_valid(ho, kh, pt, h);
315        let full_recip: T = ((kh * kw) as usize).as_().recip();
316
317        let mut htmp = vec![T::zero(); h as usize * wo];
318        unsafe {
319            for nn in 0..n as isize {
320                for cc in 0..c as isize {
321                    let in_base = nn * in_stride + cc * ic_stride;
322                    let out_base = nn * on_stride + cc * oc_stride;
323                    for y in 0..h {
324                        let row = in_base + y * ih_stride;
325                        let dst = y as usize * wo;
326                        let mut acc = T::zero();
327                        for kx in 0..kw {
328                            let ix = -pl + kx;
329                            if ix >= 0 && ix < w {
330                                acc = acc + *input_ptr.offset(row + ix);
331                            }
332                        }
333                        *htmp.get_unchecked_mut(dst) = acc;
334                        for ox in 1..wo as isize {
335                            let entering = ox - pl + kw - 1;
336                            let leaving = ox - pl - 1;
337                            if entering >= 0 && entering < w {
338                                acc = acc + *input_ptr.offset(row + entering);
339                            }
340                            if leaving >= 0 && leaving < w {
341                                acc = acc - *input_ptr.offset(row + leaving);
342                            }
343                            *htmp.get_unchecked_mut(dst + ox as usize) = acc;
344                        }
345                    }
346                    for ox in 0..wo {
347                        let mut acc = T::zero();
348                        for ky in 0..kh {
349                            let iy = -pt + ky;
350                            if iy >= 0 && iy < h {
351                                acc = acc + *htmp.get_unchecked(iy as usize * wo + ox);
352                            }
353                        }
354                        let store = |oy: usize, acc: T| {
355                            let div = if self.count_include_pad {
356                                full_recip
357                            } else {
358                                (kx_valid[ox] * ky_valid[oy]).as_().recip()
359                            };
360                            *values_ptr.offset(
361                                out_base + oy as isize * oh_stride + ox as isize * ow_stride,
362                            ) = acc * div;
363                        };
364                        store(0, acc);
365                        for oy in 1..ho as isize {
366                            let entering = oy - pt + kh - 1;
367                            let leaving = oy - pt - 1;
368                            if entering >= 0 && entering < h {
369                                acc = acc + *htmp.get_unchecked(entering as usize * wo + ox);
370                            }
371                            if leaving >= 0 && leaving < h {
372                                acc = acc - *htmp.get_unchecked(leaving as usize * wo + ox);
373                            }
374                            store(oy as usize, acc);
375                        }
376                    }
377                }
378            }
379        }
380    }
381
382    /// NHWC counterpart of `fast_2d_separable`. Channels are the innermost
383    /// (contiguous) axis, so both separable passes accumulate `C`-wide running
384    /// sums, keeping the inner channel loops contiguous. Same reassociation
385    /// caveat as the NCHW path: not bit-identical to the generic kernel.
386    unsafe fn fast_2d_separable_nhwc<T: Copy + Datum + num_traits::Float>(
387        &self,
388        input_ptr: *const T,
389        values_ptr: *mut T,
390        geo: &ConcretePoolGeometry,
391    ) where
392        usize: AsPrimitive<T>,
393    {
394        let ish = &geo.input_shape;
395        let osh = &geo.output_shape;
396        let (h, w) = (ish.hw_dims()[0] as isize, ish.hw_dims()[1] as isize);
397        let (ho, wo) = (geo.patch.output_shape[0], geo.patch.output_shape[1]);
398        let (kh, kw) =
399            (geo.patch.spec.kernel_shape[0] as isize, geo.patch.spec.kernel_shape[1] as isize);
400        let (pt, pl) = (geo.patch.pad_before[0] as isize, geo.patch.pad_before[1] as isize);
401        let ih_stride = *ish.h_stride() as isize;
402        let iw_stride = *ish.w_stride() as isize;
403        let oh_stride = *osh.h_stride() as isize;
404        let ow_stride = *osh.w_stride() as isize;
405        let n = *ish.n().unwrap_or(&1);
406        let in_stride = *ish.n_stride().unwrap_or(&0) as isize;
407        let on_stride = *osh.n_stride().unwrap_or(&0) as isize;
408        let c = *ish.c();
409
410        let axis_valid = |out: usize, k: isize, pad: isize, lim: isize| -> Vec<usize> {
411            (0..out)
412                .map(|o| {
413                    let lo = o as isize - pad;
414                    let start = (-lo).max(0);
415                    let end = (lim - lo).min(k);
416                    (end - start).max(0) as usize
417                })
418                .collect()
419        };
420        let kx_valid = axis_valid(wo, kw, pl, w);
421        let ky_valid = axis_valid(ho, kh, pt, h);
422        let full_recip: T = ((kh * kw) as usize).as_().recip();
423
424        let mut htmp = vec![T::zero(); h as usize * wo * c];
425        let mut acc = vec![T::zero(); c];
426        unsafe {
427            for nn in 0..n as isize {
428                let in_base = nn * in_stride;
429                let out_base = nn * on_stride;
430                for y in 0..h {
431                    let row = in_base + y * ih_stride;
432                    let hrow = y as usize * wo * c;
433                    acc.iter_mut().for_each(|a| *a = T::zero());
434                    for kx in 0..kw {
435                        let ix = -pl + kx;
436                        if ix >= 0 && ix < w {
437                            let p = row + ix * iw_stride;
438                            for (ch, a) in acc.iter_mut().enumerate() {
439                                *a = *a + *input_ptr.offset(p + ch as isize);
440                            }
441                        }
442                    }
443                    htmp[hrow..hrow + c].copy_from_slice(&acc);
444                    for ox in 1..wo as isize {
445                        let entering = ox - pl + kw - 1;
446                        let leaving = ox - pl - 1;
447                        if entering >= 0 && entering < w {
448                            let p = row + entering * iw_stride;
449                            for (ch, a) in acc.iter_mut().enumerate() {
450                                *a = *a + *input_ptr.offset(p + ch as isize);
451                            }
452                        }
453                        if leaving >= 0 && leaving < w {
454                            let p = row + leaving * iw_stride;
455                            for (ch, a) in acc.iter_mut().enumerate() {
456                                *a = *a - *input_ptr.offset(p + ch as isize);
457                            }
458                        }
459                        let dst = hrow + ox as usize * c;
460                        htmp[dst..dst + c].copy_from_slice(&acc);
461                    }
462                }
463                for ox in 0..wo {
464                    acc.iter_mut().for_each(|a| *a = T::zero());
465                    for ky in 0..kh {
466                        let iy = -pt + ky;
467                        if iy >= 0 && iy < h {
468                            let src = iy as usize * wo * c + ox * c;
469                            for (ch, a) in acc.iter_mut().enumerate() {
470                                *a = *a + *htmp.get_unchecked(src + ch);
471                            }
472                        }
473                    }
474                    let store = |oy: usize, acc: &[T]| {
475                        let div = if self.count_include_pad {
476                            full_recip
477                        } else {
478                            (kx_valid[ox] * ky_valid[oy]).as_().recip()
479                        };
480                        let o = out_base + oy as isize * oh_stride + ox as isize * ow_stride;
481                        for (ch, &a) in acc.iter().enumerate() {
482                            *values_ptr.offset(o + ch as isize) = a * div;
483                        }
484                    };
485                    store(0, &acc);
486                    for oy in 1..ho as isize {
487                        let entering = oy - pt + kh - 1;
488                        let leaving = oy - pt - 1;
489                        if entering >= 0 && entering < h {
490                            let src = entering as usize * wo * c + ox * c;
491                            for (ch, a) in acc.iter_mut().enumerate() {
492                                *a = *a + *htmp.get_unchecked(src + ch);
493                            }
494                        }
495                        if leaving >= 0 && leaving < h {
496                            let src = leaving as usize * wo * c + ox * c;
497                            for (ch, a) in acc.iter_mut().enumerate() {
498                                *a = *a - *htmp.get_unchecked(src + ch);
499                            }
500                        }
501                        store(oy as usize, &acc);
502                    }
503                }
504            }
505        }
506    }
507}
508
509#[cfg(test)]
510mod tests {
511    use super::*;
512    use crate::ops::cnn::PaddingSpec;
513    use crate::ops::nn::DataFormat;
514
515    fn test_case() -> (TypedModel, TVec<TValue>) {
516        let mut model = TypedModel::default();
517        let source = model.add_source("data", f32::fact([1, 3, 8, 8])).unwrap();
518        let pool_spec = PoolSpec::new(
519            DataFormat::NCHW,
520            tvec![2, 2],
521            PaddingSpec::Valid,
522            None,
523            Some(tvec![2, 2]),
524            3,
525            3,
526        );
527        let op = SumPool { pool_spec, count_include_pad: false, normalize: true };
528        let out = model.wire_node("pool", op, &[source]).unwrap();
529        model.select_output_outlets(&out).unwrap();
530        let input = ndarray::Array4::from_shape_fn((1, 3, 8, 8), |(_, c, y, x)| {
531            (c * 64 + y * 8 + x) as f32
532        })
533        .into_tensor()
534        .into_tvalue();
535        (model, tvec!(input))
536    }
537
538    #[test]
539    fn optimized_sumpool_has_concrete_geometry() {
540        let (model, input) = test_case();
541        let plain = model.clone().into_runnable().unwrap().run(input.clone()).unwrap();
542
543        let optimized = model.into_optimized().unwrap();
544        let pool = optimized
545            .nodes
546            .iter()
547            .find_map(|n| n.op_as::<OptSumPool>())
548            .expect("optimized model should contain an OptSumPool");
549        assert!(
550            pool.geometry.is_concrete(),
551            "OptSumPool geometry should be concrete after optimization"
552        );
553
554        let opt = optimized.into_runnable().unwrap().run(input).unwrap();
555        assert_eq!(*opt[0], *plain[0]);
556    }
557
558    #[test]
559    fn separable_matches_generic_kernel() {
560        let (c, h, w) = (5usize, 7usize, 9usize);
561        let pool_spec = PoolSpec::new(
562            DataFormat::NCHW,
563            tvec![3, 3],
564            PaddingSpec::SameUpper,
565            None,
566            Some(tvec![1, 1]),
567            c,
568            c,
569        );
570        let op = OptSumPool {
571            pool_spec: pool_spec.clone(),
572            count_include_pad: false,
573            normalize: true,
574            geometry: pool_spec
575                .compute_geo(&[1.to_dim(), c.to_dim(), h.to_dim(), w.to_dim()])
576                .unwrap(),
577        };
578        let input: Tensor = ndarray::Array4::from_shape_fn((1, c, h, w), |(_, cc, y, x)| {
579            ((cc * 17 + y * 3 + x) % 13) as f32 - 6.0
580        })
581        .into_tensor();
582
583        // generic zoned kernel (knob off by default)
584        let generic = op.eval(tvec![input.clone().into_tvalue()]).unwrap();
585        let generic = generic[0].try_as_plain().unwrap().as_slice::<f32>().unwrap().to_vec();
586
587        // separable kernel, called directly
588        let geo = op.geometry.to_concrete(input.shape()).unwrap();
589        let mut out = Tensor::zero::<f32>(&geo.output_shape.shape).unwrap();
590        unsafe {
591            op.fast_2d_separable::<f32>(
592                input.as_ptr::<f32>().unwrap(),
593                out.as_ptr_mut::<f32>().unwrap(),
594                geo.as_ref(),
595            );
596        }
597        let sep = out.try_as_plain().unwrap().as_slice::<f32>().unwrap();
598
599        let max_abs = generic.iter().zip(sep).map(|(a, b)| (a - b).abs()).fold(0f32, f32::max);
600        assert!(max_abs < 1e-4, "separable vs generic max abs diff {max_abs}");
601    }
602
603    #[test]
604    fn separable_nhwc_matches_generic_kernel() {
605        let (c, h, w) = (5usize, 7usize, 9usize);
606        let pool_spec = PoolSpec::new(
607            DataFormat::NHWC,
608            tvec![3, 3],
609            PaddingSpec::SameUpper,
610            None,
611            Some(tvec![1, 1]),
612            c,
613            c,
614        );
615        let op = OptSumPool {
616            pool_spec: pool_spec.clone(),
617            count_include_pad: false,
618            normalize: true,
619            geometry: pool_spec
620                .compute_geo(&[1.to_dim(), h.to_dim(), w.to_dim(), c.to_dim()])
621                .unwrap(),
622        };
623        let input: Tensor = ndarray::Array4::from_shape_fn((1, h, w, c), |(_, y, x, cc)| {
624            ((cc * 17 + y * 3 + x) % 13) as f32 - 6.0
625        })
626        .into_tensor();
627
628        // generic zoned kernel (knob off by default)
629        let generic = op.eval(tvec![input.clone().into_tvalue()]).unwrap();
630        let generic = generic[0].try_as_plain().unwrap().as_slice::<f32>().unwrap().to_vec();
631
632        // separable NHWC kernel, called directly
633        let geo = op.geometry.to_concrete(input.shape()).unwrap();
634        let mut out = Tensor::zero::<f32>(&geo.output_shape.shape).unwrap();
635        unsafe {
636            op.fast_2d_separable_nhwc::<f32>(
637                input.as_ptr::<f32>().unwrap(),
638                out.as_ptr_mut::<f32>().unwrap(),
639                geo.as_ref(),
640            );
641        }
642        let sep = out.try_as_plain().unwrap().as_slice::<f32>().unwrap();
643
644        let max_abs = generic.iter().zip(sep).map(|(a, b)| (a - b).abs()).fold(0f32, f32::max);
645        assert!(max_abs < 1e-4, "separable NHWC vs generic max abs diff {max_abs}");
646    }
647}