cubek-reduce 0.3.0-pre.2

use super::{
    All, Any, Max, MaxAbs, Mean, Min, Prod, ReduceFamily, ReduceInstruction, ReduceRequirements,
    SharedAccumulator, Sum,
};
use crate::components::instructions::{
    Accumulator, AccumulatorFormat, Item, ReduceOutputMode, SharedAccumulatorKind, TopK,
};
use crate::{
    ReduceDtypes,
    components::{
        instructions::{ReduceStep, Value},
        precision::ReducePrecision,
    },
};
use cubecl::{
    ir::{ElemType, FloatKind, IntKind, UIntKind},
    prelude::*,
};
use serde::{Deserialize, Serialize};

#[derive(Debug, CubeType, Clone)]
pub enum ReduceOperation {
    Sum(Sum),
    Prod(Prod),
    Mean(Mean),
    MaxAbs(MaxAbs),
    Max(Max),
    Min(Min),
    TopK(TopK),
    Any(Any),
    All(All),
}

#[derive_cube_comptime]
#[derive(Serialize, Deserialize)]
pub enum ReduceOperationConfig {
    Sum,
    Prod,
    Mean,
    MaxAbs,
    ArgMax,
    ArgMin,
    Max,
    Min,
    ArgTopK(usize),
    TopK(usize),
    Any,
    All,
}

impl ReduceOperationConfig {
    /// Shared-memory bytes one accumulator slot uses (total usage is this times the
    /// slot count). `acc_elem_size` is the accumulation element size (`P::EA`),
    /// `vector_size` the input vectorization. Mirrors each instruction's
    /// `SharedAccumulator` layout: one value slice, plus a `u32` index slice for
    /// `Arg*`, scaled by `k` for top-k.
    pub fn shared_memory_bytes_per_accumulator(
        &self,
        acc_elem_size: usize,
        vector_size: usize,
    ) -> usize {
        // Index slices are `Vector<u32, SI>` for every instruction that has them
        // (`ArgAccumulator` and `TopKSharedAccumulator`), so the index element is
        // always u32. Revisit this if indices ever widen (e.g. u64 coordinates).
        let index_elem_size = core::mem::size_of::<u32>();
        let (value_slices, index_slices) = match self {
            ReduceOperationConfig::Sum
            | ReduceOperationConfig::Prod
            | ReduceOperationConfig::Mean
            | ReduceOperationConfig::MaxAbs
            | ReduceOperationConfig::Max
            | ReduceOperationConfig::Min
            | ReduceOperationConfig::Any
            | ReduceOperationConfig::All => (1, 0),
            ReduceOperationConfig::ArgMax | ReduceOperationConfig::ArgMin => (1, 1),
            ReduceOperationConfig::ArgTopK(k) => (*k, *k),
            ReduceOperationConfig::TopK(k) => (*k, 0),
        };
        (value_slices * acc_elem_size + index_slices * index_elem_size) * vector_size
    }

    /// Computes the best case precision for the given config.
    pub fn precision(&self, input: ElemType, output: Option<ElemType>) -> ReduceDtypes {
        match self {
            ReduceOperationConfig::Sum
            | ReduceOperationConfig::Prod
            | ReduceOperationConfig::Mean => {}
            // No benefit to mixed precision accumulation.
            ReduceOperationConfig::MaxAbs
            | ReduceOperationConfig::Max
            | ReduceOperationConfig::TopK(_)
            | ReduceOperationConfig::Min => {
                return ReduceDtypes {
                    input: input.into(),
                    output: input.into(),
                    accumulation: input.into(),
                };
            }
            // The output is not a value of the input but an index (Arg*) or a
            // logical flag (`Any` / `All`), so the caller must pick its dtype;
            // the conversion happens for free in the final output write. The
            // accumulator stays narrow (= input): indices live in a separate
            // u32 accumulator, and logical flags only ever hold 0/1.
            ReduceOperationConfig::ArgMax
            | ReduceOperationConfig::ArgMin
            | ReduceOperationConfig::ArgTopK(_)
            | ReduceOperationConfig::Any
            | ReduceOperationConfig::All => {
                return ReduceDtypes {
                    input: input.into(),
                    output: output
                        .expect("ArgMax, ArgMin, ArgTopK, Any and All must specify output type")
                        .into(),
                    accumulation: input.into(),
                };
            }
        };

        match input {
            ElemType::Float(kind) => {
                let acc = match kind {
                    FloatKind::F64 => f64::as_type_native_unchecked(),
                    _ => f32::as_type_native_unchecked(),
                };

                ReduceDtypes {
                    input: input.into(),
                    output: input.into(),
                    accumulation: acc.storage_type(),
                }
            }
            ElemType::Int(kind) => {
                let acc = match kind {
                    IntKind::I64 => i64::as_type_native_unchecked(),
                    _ => i32::as_type_native_unchecked(),
                };

                ReduceDtypes {
                    input: input.into(),
                    output: input.into(),
                    accumulation: acc.storage_type(),
                }
            }
            ElemType::UInt(kind) => {
                let acc = match kind {
                    UIntKind::U64 => u64::as_type_native_unchecked(),
                    _ => u32::as_type_native_unchecked(),
                };

                ReduceDtypes {
                    input: input.into(),
                    output: input.into(),
                    accumulation: acc.storage_type(),
                }
            }
            ElemType::Bool => panic!("Can't reduce on booleans"),
        }
    }
}

impl ReduceFamily for ReduceOperation {
    type Instruction<P: ReducePrecision> = Self;
    type Config = ReduceOperationConfig;
}

#[derive(CubeType)]
pub struct DynamicSharedAccumulator<P: ReducePrecision> {
    pub elements: SharedAccumulatorKind<Vector<P::EA, P::SI>>,
    pub args: SharedAccumulatorKind<Vector<u32, P::SI>>,
}

#[derive(CubeType)]
pub struct DynamicAccumulator<P: ReducePrecision> {
    pub elements: Value<Vector<P::EA, P::SI>>,
    pub args: Value<Vector<u32, P::SI>>,
}

#[cube]
impl<P: ReducePrecision, I: ReduceInstruction<P>> SharedAccumulator<P, I>
    for DynamicSharedAccumulator<P>
{
    fn allocate(#[comptime] length: usize, #[comptime] coordinate: bool, inst: &I) -> Self {
        let format = I::accumulator_format(inst);
        match comptime!(format) {
            AccumulatorFormat::Single => {
                let elements = Shared::new_slice(length);
                // TODO how to put multiple?
                let args = if coordinate {
                    let args = Shared::new_slice(length);
                    SharedAccumulatorKind::new_Single(args)
                } else {
                    SharedAccumulatorKind::new_None()
                };
                DynamicSharedAccumulator::<P> {
                    elements: SharedAccumulatorKind::new_Single(elements),
                    args,
                }
            }
            AccumulatorFormat::Multiple(len) => {
                let mut elements = Sequence::new();
                #[unroll]
                for _ in 0..len {
                    elements.push(Shared::new_slice(length));
                }

                if comptime!(!coordinate) {
                    DynamicSharedAccumulator::<P> {
                        elements: SharedAccumulatorKind::new_Multiple(elements),
                        args: SharedAccumulatorKind::new_None(),
                    }
                } else {
                    let mut args = Sequence::new();
                    #[unroll]
                    for _ in 0..len {
                        args.push(Shared::new_slice(length));
                    }
                    DynamicSharedAccumulator::<P> {
                        elements: SharedAccumulatorKind::new_Multiple(elements),
                        args: SharedAccumulatorKind::new_Multiple(args),
                    }
                }
            }
        }
    }

    fn read(accumulator: &Self, index: usize) -> Accumulator<P> {
        let elements = accumulator.elements.get(index);
        let args = accumulator.args.get(index);

        Accumulator::<P> { elements, args }
    }

    fn write(accumulator: &mut Self, index: usize, item: Accumulator<P>) {
        accumulator.elements.set(index, item.elements);
        accumulator.args.set(index, item.args);
    }
}

#[cube]
impl<P: ReducePrecision> ReduceInstruction<P> for ReduceOperation {
    type SharedAccumulator = DynamicSharedAccumulator<P>;
    type Config = ReduceOperationConfig;

    fn requirements(this: &Self) -> ReduceRequirements {
        match this {
            ReduceOperation::Sum(sum) => <Sum as ReduceInstruction<P>>::requirements(sum),
            ReduceOperation::Prod(prod) => <Prod as ReduceInstruction<P>>::requirements(prod),
            ReduceOperation::Mean(mean) => <Mean as ReduceInstruction<P>>::requirements(mean),
            ReduceOperation::MaxAbs(max_abs) => {
                <MaxAbs as ReduceInstruction<P>>::requirements(max_abs)
            }
            ReduceOperation::TopK(topk) => <TopK as ReduceInstruction<P>>::requirements(topk),
            ReduceOperation::Max(max) => <Max as ReduceInstruction<P>>::requirements(max),
            ReduceOperation::Min(min) => <Min as ReduceInstruction<P>>::requirements(min),
            ReduceOperation::Any(any) => <Any as ReduceInstruction<P>>::requirements(any),
            ReduceOperation::All(all) => <All as ReduceInstruction<P>>::requirements(all),
        }
    }

    fn accumulator_format(this: &Self) -> comptime_type!(AccumulatorFormat) {
        match this {
            ReduceOperation::Sum(sum) => <Sum as ReduceInstruction<P>>::accumulator_format(sum),
            ReduceOperation::Prod(prod) => <Prod as ReduceInstruction<P>>::accumulator_format(prod),
            ReduceOperation::Mean(mean) => <Mean as ReduceInstruction<P>>::accumulator_format(mean),
            ReduceOperation::MaxAbs(maxabs) => {
                <MaxAbs as ReduceInstruction<P>>::accumulator_format(maxabs)
            }
            ReduceOperation::Max(max) => <Max as ReduceInstruction<P>>::accumulator_format(max),
            ReduceOperation::Min(min) => <Min as ReduceInstruction<P>>::accumulator_format(min),
            ReduceOperation::TopK(topk) => <TopK as ReduceInstruction<P>>::accumulator_format(topk),
            ReduceOperation::Any(any) => <Any as ReduceInstruction<P>>::accumulator_format(any),
            ReduceOperation::All(all) => <All as ReduceInstruction<P>>::accumulator_format(all),
        }
    }

    fn from_config(#[comptime] config: Self::Config) -> Self {
        match config {
            ReduceOperationConfig::Sum => ReduceOperation::new_Sum(Sum {}),
            ReduceOperationConfig::Prod => ReduceOperation::new_Prod(Prod {}),
            ReduceOperationConfig::Mean => ReduceOperation::new_Mean(Mean { sum: Sum {} }),
            ReduceOperationConfig::MaxAbs => ReduceOperation::new_MaxAbs(MaxAbs {}),
            ReduceOperationConfig::ArgMax => ReduceOperation::new_Max(Max {
                output: ReduceOutputMode::Indices,
            }),
            ReduceOperationConfig::ArgMin => ReduceOperation::new_Min(Min {
                output: ReduceOutputMode::Indices,
            }),
            ReduceOperationConfig::ArgTopK(k) => ReduceOperation::new_TopK(TopK {
                k,
                output: ReduceOutputMode::Indices,
            }),
            ReduceOperationConfig::Max => ReduceOperation::new_Max(Max {
                output: ReduceOutputMode::Values,
            }),
            ReduceOperationConfig::Min => ReduceOperation::new_Min(Min {
                output: ReduceOutputMode::Values,
            }),
            ReduceOperationConfig::TopK(k) => ReduceOperation::new_TopK(TopK {
                k,
                output: ReduceOutputMode::Values,
            }),
            ReduceOperationConfig::Any => ReduceOperation::new_Any(Any {}),
            ReduceOperationConfig::All => ReduceOperation::new_All(All {}),
        }
    }

    fn null_input(this: &Self) -> Vector<P::EI, P::SI> {
        match this {
            ReduceOperation::Sum(sum) => <Sum as ReduceInstruction<P>>::null_input(sum),
            ReduceOperation::Prod(prod) => <Prod as ReduceInstruction<P>>::null_input(prod),
            ReduceOperation::Mean(mean) => <Mean as ReduceInstruction<P>>::null_input(mean),
            ReduceOperation::MaxAbs(maxabs) => <MaxAbs as ReduceInstruction<P>>::null_input(maxabs),
            ReduceOperation::Max(max) => <Max as ReduceInstruction<P>>::null_input(max),
            ReduceOperation::Min(min) => <Min as ReduceInstruction<P>>::null_input(min),
            ReduceOperation::TopK(topk) => <TopK as ReduceInstruction<P>>::null_input(topk),
            ReduceOperation::Any(any) => <Any as ReduceInstruction<P>>::null_input(any),
            ReduceOperation::All(all) => <All as ReduceInstruction<P>>::null_input(all),
        }
    }

    fn null_accumulator(this: &Self) -> Accumulator<P> {
        match this {
            ReduceOperation::Sum(sum) => <Sum as ReduceInstruction<P>>::null_accumulator(sum),
            ReduceOperation::Mean(sum) => <Mean as ReduceInstruction<P>>::null_accumulator(sum),
            ReduceOperation::Prod(prod) => <Prod as ReduceInstruction<P>>::null_accumulator(prod),
            ReduceOperation::MaxAbs(maxabs) => {
                <MaxAbs as ReduceInstruction<P>>::null_accumulator(maxabs)
            }
            ReduceOperation::Max(max) => <Max as ReduceInstruction<P>>::null_accumulator(max),
            ReduceOperation::Min(min) => <Min as ReduceInstruction<P>>::null_accumulator(min),
            ReduceOperation::TopK(topk) => <TopK as ReduceInstruction<P>>::null_accumulator(topk),
            ReduceOperation::Any(any) => <Any as ReduceInstruction<P>>::null_accumulator(any),
            ReduceOperation::All(all) => <All as ReduceInstruction<P>>::null_accumulator(all),
        }
    }

    fn reduce(
        this: &Self,
        accumulator: &mut Accumulator<P>,
        item: Item<P>,
        #[comptime] reduce_step: ReduceStep,
    ) {
        match this {
            ReduceOperation::Sum(sum) => {
                <Sum as ReduceInstruction<P>>::reduce(sum, accumulator, item, reduce_step)
            }
            ReduceOperation::Prod(sum) => {
                <Prod as ReduceInstruction<P>>::reduce(sum, accumulator, item, reduce_step)
            }
            ReduceOperation::Mean(sum) => {
                <Mean as ReduceInstruction<P>>::reduce(sum, accumulator, item, reduce_step)
            }
            ReduceOperation::MaxAbs(maxabs) => {
                <MaxAbs as ReduceInstruction<P>>::reduce(maxabs, accumulator, item, reduce_step)
            }
            ReduceOperation::Max(max) => {
                <Max as ReduceInstruction<P>>::reduce(max, accumulator, item, reduce_step)
            }
            ReduceOperation::Min(min) => {
                <Min as ReduceInstruction<P>>::reduce(min, accumulator, item, reduce_step)
            }
            ReduceOperation::TopK(topk) => {
                <TopK as ReduceInstruction<P>>::reduce(topk, accumulator, item, reduce_step)
            }
            ReduceOperation::Any(any) => {
                <Any as ReduceInstruction<P>>::reduce(any, accumulator, item, reduce_step)
            }
            ReduceOperation::All(all) => {
                <All as ReduceInstruction<P>>::reduce(all, accumulator, item, reduce_step)
            }
        }
    }

    fn plane_reduce_inplace(this: &Self, accumulator: &mut Accumulator<P>) {
        match this {
            ReduceOperation::Sum(sum) => {
                <Sum as ReduceInstruction<P>>::plane_reduce_inplace(sum, accumulator)
            }
            ReduceOperation::Prod(prod) => {
                <Prod as ReduceInstruction<P>>::plane_reduce_inplace(prod, accumulator)
            }
            ReduceOperation::Mean(mean) => {
                <Mean as ReduceInstruction<P>>::plane_reduce_inplace(mean, accumulator)
            }
            ReduceOperation::MaxAbs(max_abs) => {
                <MaxAbs as ReduceInstruction<P>>::plane_reduce_inplace(max_abs, accumulator)
            }
            ReduceOperation::Max(max) => {
                <Max as ReduceInstruction<P>>::plane_reduce_inplace(max, accumulator)
            }
            ReduceOperation::Min(min) => {
                <Min as ReduceInstruction<P>>::plane_reduce_inplace(min, accumulator)
            }
            ReduceOperation::TopK(topk) => {
                <TopK as ReduceInstruction<P>>::plane_reduce_inplace(topk, accumulator)
            }
            ReduceOperation::Any(any) => {
                <Any as ReduceInstruction<P>>::plane_reduce_inplace(any, accumulator)
            }
            ReduceOperation::All(all) => {
                <All as ReduceInstruction<P>>::plane_reduce_inplace(all, accumulator)
            }
        }
    }

    fn fuse_accumulators(this: &Self, accumulator: &mut Accumulator<P>, other: &Accumulator<P>) {
        match this {
            ReduceOperation::Sum(sum) => {
                <Sum as ReduceInstruction<P>>::fuse_accumulators(sum, accumulator, other)
            }
            ReduceOperation::Prod(prod) => {
                <Prod as ReduceInstruction<P>>::fuse_accumulators(prod, accumulator, other)
            }
            ReduceOperation::Mean(mean) => {
                <Mean as ReduceInstruction<P>>::fuse_accumulators(mean, accumulator, other)
            }
            ReduceOperation::MaxAbs(maxabs) => {
                <MaxAbs as ReduceInstruction<P>>::fuse_accumulators(maxabs, accumulator, other)
            }
            ReduceOperation::Max(max) => {
                <Max as ReduceInstruction<P>>::fuse_accumulators(max, accumulator, other)
            }
            ReduceOperation::Min(min) => {
                <Min as ReduceInstruction<P>>::fuse_accumulators(min, accumulator, other)
            }
            ReduceOperation::TopK(topk) => {
                <TopK as ReduceInstruction<P>>::fuse_accumulators(topk, accumulator, other)
            }
            ReduceOperation::Any(any) => {
                <Any as ReduceInstruction<P>>::fuse_accumulators(any, accumulator, other)
            }
            ReduceOperation::All(all) => {
                <All as ReduceInstruction<P>>::fuse_accumulators(all, accumulator, other)
            }
        }
    }

    fn output_mode(this: &Self) -> comptime_type!(ReduceOutputMode) {
        match this {
            ReduceOperation::Sum(sum) => <Sum as ReduceInstruction<P>>::output_mode(sum),
            ReduceOperation::Prod(prod) => <Prod as ReduceInstruction<P>>::output_mode(prod),
            ReduceOperation::Mean(mean) => <Mean as ReduceInstruction<P>>::output_mode(mean),
            ReduceOperation::MaxAbs(maxabs) => {
                <MaxAbs as ReduceInstruction<P>>::output_mode(maxabs)
            }
            ReduceOperation::Max(max) => <Max as ReduceInstruction<P>>::output_mode(max),
            ReduceOperation::Min(min) => <Min as ReduceInstruction<P>>::output_mode(min),
            ReduceOperation::TopK(topk) => <TopK as ReduceInstruction<P>>::output_mode(topk),
            ReduceOperation::Any(any) => <Any as ReduceInstruction<P>>::output_mode(any),
            ReduceOperation::All(all) => <All as ReduceInstruction<P>>::output_mode(all),
        }
    }

    fn to_output_parallel<Out: Numeric, Idx: Numeric>(
        this: &Self,
        accumulator: Accumulator<P>,
        shape_axis_reduce: usize,
    ) -> (Value<Out>, Value<Idx>) {
        match this {
            ReduceOperation::Sum(sum) => <Sum as ReduceInstruction<P>>::to_output_parallel::<
                Out,
                Idx,
            >(sum, accumulator, shape_axis_reduce),
            ReduceOperation::Prod(prod) => <Prod as ReduceInstruction<P>>::to_output_parallel::<
                Out,
                Idx,
            >(prod, accumulator, shape_axis_reduce),
            ReduceOperation::Mean(mean) => <Mean as ReduceInstruction<P>>::to_output_parallel::<
                Out,
                Idx,
            >(mean, accumulator, shape_axis_reduce),
            ReduceOperation::MaxAbs(maxabs) => {
                <MaxAbs as ReduceInstruction<P>>::to_output_parallel::<Out, Idx>(
                    maxabs,
                    accumulator,
                    shape_axis_reduce,
                )
            }
            ReduceOperation::Max(max) => <Max as ReduceInstruction<P>>::to_output_parallel::<
                Out,
                Idx,
            >(max, accumulator, shape_axis_reduce),
            ReduceOperation::Min(min) => <Min as ReduceInstruction<P>>::to_output_parallel::<
                Out,
                Idx,
            >(min, accumulator, shape_axis_reduce),
            ReduceOperation::TopK(topk) => <TopK as ReduceInstruction<P>>::to_output_parallel::<
                Out,
                Idx,
            >(topk, accumulator, shape_axis_reduce),
            ReduceOperation::Any(any) => <Any as ReduceInstruction<P>>::to_output_parallel::<
                Out,
                Idx,
            >(any, accumulator, shape_axis_reduce),
            ReduceOperation::All(all) => <All as ReduceInstruction<P>>::to_output_parallel::<
                Out,
                Idx,
            >(all, accumulator, shape_axis_reduce),
        }
    }

    fn to_output_perpendicular<Out: Numeric, Idx: Numeric>(
        this: &Self,
        accumulator: Accumulator<P>,
        shape_axis_reduce: usize,
    ) -> (Value<Vector<Out, P::SI>>, Value<Vector<Idx, P::SI>>) {
        match this {
            ReduceOperation::Sum(sum) => <Sum as ReduceInstruction<P>>::to_output_perpendicular::<
                Out,
                Idx,
            >(sum, accumulator, shape_axis_reduce),
            ReduceOperation::Prod(prod) => {
                <Prod as ReduceInstruction<P>>::to_output_perpendicular::<Out, Idx>(
                    prod,
                    accumulator,
                    shape_axis_reduce,
                )
            }
            ReduceOperation::Mean(mean) => {
                <Mean as ReduceInstruction<P>>::to_output_perpendicular::<Out, Idx>(
                    mean,
                    accumulator,
                    shape_axis_reduce,
                )
            }
            ReduceOperation::MaxAbs(maxabs) => {
                <MaxAbs as ReduceInstruction<P>>::to_output_perpendicular::<Out, Idx>(
                    maxabs,
                    accumulator,
                    shape_axis_reduce,
                )
            }
            ReduceOperation::Max(max) => <Max as ReduceInstruction<P>>::to_output_perpendicular::<
                Out,
                Idx,
            >(max, accumulator, shape_axis_reduce),
            ReduceOperation::Min(min) => <Min as ReduceInstruction<P>>::to_output_perpendicular::<
                Out,
                Idx,
            >(min, accumulator, shape_axis_reduce),
            ReduceOperation::TopK(topk) => {
                <TopK as ReduceInstruction<P>>::to_output_perpendicular::<Out, Idx>(
                    topk,
                    accumulator,
                    shape_axis_reduce,
                )
            }
            ReduceOperation::Any(any) => <Any as ReduceInstruction<P>>::to_output_perpendicular::<
                Out,
                Idx,
            >(any, accumulator, shape_axis_reduce),
            ReduceOperation::All(all) => <All as ReduceInstruction<P>>::to_output_perpendicular::<
                Out,
                Idx,
            >(all, accumulator, shape_axis_reduce),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The key benefit of `Any` / `All` over the sum-based emulation is that the
    /// dynamic dispatch path keeps `accumulation = input`: the accumulator only
    /// ever holds 0/1 flags, so there is no need to widen (and no overflow) the
    /// way `Sum` / `Mean` do. The output dtype is the caller-provided flag
    /// storage (e.g. the u8/u32 backing of a bool tensor), like Arg* indices.
    /// Check both for a narrow float (f16, the type that motivated this work)
    /// and an integer input.
    /// Pin the footprint for one op of each layout. f32 is 4 bytes, so `ArgTopK(k)`
    /// is `8 * k` bytes.
    #[test]
    fn shared_memory_footprint_matches_accumulator_layout() {
        use ReduceOperationConfig::*;
        let cases = [
            (Sum, 4),
            (ArgMax, 8),
            (TopK(7), 7 * 4),
            (ArgTopK(13), 13 * 8),
        ];
        for (config, expected) in cases {
            assert_eq!(
                config.shared_memory_bytes_per_accumulator(4, 1),
                expected,
                "footprint for {config:?}"
            );
        }
        // Vectorization scales every slice uniformly.
        assert_eq!(
            ArgTopK(13).shared_memory_bytes_per_accumulator(4, 4),
            13 * 8 * 4
        );
    }

    #[test]
    fn any_all_precision_keeps_accumulation_narrow() {
        let inputs = [ElemType::Float(FloatKind::F16), ElemType::Int(IntKind::I32)];
        let output = ElemType::UInt(UIntKind::U8);
        for config in [ReduceOperationConfig::Any, ReduceOperationConfig::All] {
            for input in inputs {
                let dtypes = config.precision(input, Some(output));
                let expected: StorageType = input.into();
                assert_eq!(dtypes.input, expected, "input for {input:?}");
                assert_eq!(
                    dtypes.output,
                    output.into(),
                    "output must follow the requested flag storage for {input:?}"
                );
                assert_eq!(
                    dtypes.accumulation, expected,
                    "accumulation must stay narrow for {input:?}"
                );
            }
        }
    }
}