tract_core/ops/nn/softmax/

mod.rs

mod fixedpoint;
pub mod math;

use math::{
    convert_scale_to_mult_shift, exp_on_negative_values, get_reciprocal, rescale,
    rounding_divide_by_pot, saturating_rounding_doubling_high_mul,
    saturating_rounding_multiply_by_pot,
};
use num_traits::Float;
use std::fmt::Debug;
use tract_num_traits::Zero;

use crate::internal::*;
use ndarray::prelude::*;
use tract_linalg::routines::Func;

#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq, Default)]
pub enum SoftmaxKind {
    #[default]
    Softmax,
    LogSoftmax,
}

#[derive(Debug, Clone, new, Hash, Default, PartialEq, Eq)]
pub struct Softmax {
    pub axes: TVec<usize>,
    pub quant_output_dt: Option<DatumType>,
    pub kind: SoftmaxKind,
}

impl Op for Softmax {
    fn name(&self) -> StaticName {
        match self.kind {
            SoftmaxKind::Softmax => "Softmax".into(),
            SoftmaxKind::LogSoftmax => "LogSoftmax".into(),
        }
    }

    fn info(&self) -> TractResult<Vec<String>> {
        Ok(vec![format!("Axis: {:?}", self.axes)])
    }

    op_as_typed_op!();
}

impl TypedOp for Softmax {
    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
        let dt = inputs[0].datum_type;
        if dt.is_float() {
            ensure!(
                self.quant_output_dt.is_none(),
                "Float softmax should not have quant_output_dt, have {:?}",
                self.quant_output_dt
            );
        } else if dt.is_quantized() {
            ensure!(
                self.quant_output_dt.map(|q| q.is_quantized()).unwrap_or(false),
                "Quantized softmax should have a quantized output type (got {:?})",
                self.quant_output_dt
            );
        } else {
            bail!(
                "Unsupported datum type in softmax: input type {:?}, output type {:?}",
                dt,
                self.quant_output_dt
            );
        }

        let fact = self.quant_output_dt.unwrap_or(dt).fact(inputs[0].shape.clone());
        Ok(tvec!(fact))
    }

    fn input_roi(
        &self,
        model: &TypedModel,
        node: &TypedNode,
    ) -> TractResult<Option<TVec<Option<TDim>>>> {
        crate::optim::propagate_roi::bubble_roi(model, node)
    }

    fn axes_mapping(
        &self,
        inputs: &[&TypedFact],
        outputs: &[&TypedFact],
    ) -> TractResult<AxesMapping> {
        AxesMapping::natural(inputs, outputs)
    }

    fn change_axes(
        &self,
        model: &TypedModel,
        node: &TypedNode,
        _io: InOut,
        change: &AxisOp,
    ) -> TractResult<Option<AxisChangeConsequence>> {
        let axes: Option<TVec<usize>> =
            self.axes.iter().map(|it| change.transform_axis(*it)).collect();
        if let Some(axes) = axes {
            Ok(Some(AxisChangeConsequence::new(
                model,
                node,
                Some(Box::new(Softmax { axes, ..self.clone() })),
                change,
            )))
        } else {
            Ok(None)
        }
    }

    as_op!();
}

impl EvalOp for Softmax {
    op_out_of_plan!();

    fn eval(&self, _ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
        let input = args_1!(inputs);
        let dt = input.datum_type();

        let output = match dt {
            DatumType::F64 => self.eval_t::<f64>(input)?,
            DatumType::F32 => self.eval_t::<f32>(input)?,
            DatumType::F16 => self.eval_t::<f16>(input)?,
            DatumType::QI8(_) | DatumType::QU8(_) => self.eval_quant(input)?,
            dt => bail!("Unsupported type {dt:?}"),
        };
        Ok(output)
    }
}

impl Softmax {
    fn eval_t<T>(&self, input: TValue) -> TractResult<TVec<TValue>>
    where
        T: Float + Datum + std::iter::Sum,
    {
        let mut output = input.into_tensor();

        // Fast path: when the softmax axes are the trailing axes of a
        // C-contiguous tensor, the buffer is [n_rows, row_len] row-major
        // (row_len = product of the softmax dims), so every softmax region is
        // one contiguous row. Rows split across threads while each row's
        // reduction stays serial, keeping the result bit-identical.
        let rank = output.rank();
        let mut sm_axes: TVec<usize> = self.axes.clone();
        sm_axes.sort_unstable();
        let trailing = !sm_axes.is_empty()
            && sm_axes.len() <= rank
            && sm_axes.iter().enumerate().all(|(i, &a)| a == rank - sm_axes.len() + i);
        if trailing
            && matches!(T::datum_type(), DatumType::F32 | DatumType::F16)
            && output.strides() == &*Tensor::natural_strides(output.shape())
        {
            let row_len: usize = sm_axes.iter().map(|&a| output.shape()[a]).product();
            if row_len > 0 {
                let kind = self.kind;
                let mut output_plain = output.try_as_plain_ram_mut()?;
                // Dtype-concrete `as_slice_mut` (checked, no transmute) since T is
                // f32/f16 here; keeps this core path free of new `unsafe`.
                if T::datum_type() == f32::datum_type()
                    && kind == SoftmaxKind::Softmax
                    && row_len.is_multiple_of(4)
                    && let Some(rows_kernel) = tract_linalg::routines::softmax_rows_f32()
                {
                    let data = output_plain.as_slice_mut::<f32>()?;
                    let total = data.len();
                    tract_linalg::multithread::par_chunks_mut(data, row_len, total, |_, chunk| {
                        rows_kernel(chunk, row_len);
                        Ok(())
                    })?;
                } else if T::datum_type() == f32::datum_type() {
                    let data = output_plain.as_slice_mut::<f32>()?;
                    let total = data.len();
                    let kernels = F32Kernels::new()?;
                    tract_linalg::multithread::par_chunks_mut(data, row_len, total, |_, chunk| {
                        for row in chunk.chunks_mut(row_len) {
                            softmax_inner_slice_f32(row, kind, &kernels)?;
                        }
                        Ok(())
                    })?;
                } else {
                    let data = output_plain.as_slice_mut::<f16>()?;
                    let total = data.len();
                    tract_linalg::multithread::par_chunks_mut(data, row_len, total, |_, chunk| {
                        for row in chunk.chunks_mut(row_len) {
                            self.softmax_inner_slice_f16(row, kind)?;
                        }
                        Ok(())
                    })?;
                }
                return Ok(tvec!(output.into_tvalue()));
            }
        }

        // Slow path: strided or non-trailing axes, iterated with ndarray view
        // collapsing (the contiguous-row fast path above did not apply).
        let mut iterating_shape: TVec<usize> = output.shape().into();

        for i in 0..iterating_shape.len() {
            if self.axes.contains(&i) {
                iterating_shape[i] = 1
            }
        }

        let mut output_plain = output.try_as_plain_ram_mut()?;
        let mut view = output_plain.to_array_view_mut::<T>()?;
        let kernels =
            if T::datum_type() == f32::datum_type() { Some(F32Kernels::new()?) } else { None };

        for it_coords in tract_ndarray::indices(&*iterating_shape) {
            let mut view = view.view_mut();
            for ix in 0..iterating_shape.len() {
                if !self.axes.contains(&ix) {
                    view.collapse_axis(Axis(ix), it_coords[ix]);
                }
            }
            if let Some(slice) =
                view.as_slice_mut().filter(|_| T::datum_type() == f32::datum_type())
            {
                let slice: &mut [f32] = unsafe { std::mem::transmute(slice) };
                softmax_inner_slice_f32(
                    slice,
                    self.kind,
                    kernels.as_ref().context("f32 kernels")?,
                )?;
            } else if let Some(slice) =
                view.as_slice_mut().filter(|_| T::datum_type() == f16::datum_type())
            {
                let slice: &mut [f16] = unsafe { std::mem::transmute(slice) };
                self.softmax_inner_slice_f16(slice, self.kind)?;
            } else {
                softmax_inner(view, self.kind);
            }
        }

        Ok(tvec!(output.into_tvalue()))
    }

    fn eval_quant(&self, input: TValue) -> TractResult<TVec<TValue>> {
        if self.kind == SoftmaxKind::LogSoftmax {
            bail!("Quantized LogSoftmax is not supported")
        }
        let mut iterating_shape: TVec<usize> = input.shape().into();
        let output_dt =
            self.quant_output_dt.context("Quandized softmax eval with no output type")?;

        for i in 0..iterating_shape.len() {
            if self.axes.contains(&i) {
                iterating_shape[i] = 1
            }
        }

        // All operations will be done in u8, we will cast the result appropriately afterward.
        let src_is_signed = input.datum_type().is_signed();
        let out_is_signed = output_dt.is_signed();
        let in_qp = input.datum_type().qparams().unwrap(); // Checked as we are in the quant case
        let out_qp = output_dt.qparams().unwrap(); // Checked as we are in the quant case
        let mut output = unsafe { input.into_tensor().into_array_unchecked::<u8>() };

        for it_coords in tract_ndarray::indices(&*iterating_shape) {
            let mut view = output.view_mut();
            for ix in 0..iterating_shape.len() {
                if !self.axes.contains(&ix) {
                    view.collapse_axis(Axis(ix), it_coords[ix]);
                }
            }
            softmax_quant_inner(view, src_is_signed, in_qp, out_is_signed, out_qp);
        }

        let mut output_tensor = output.into_tensor();
        unsafe { output_tensor.set_datum_type(output_dt) };
        Ok(tvec!(output_tensor.into_tvalue()))
    }

    fn softmax_inner_slice_f16(&self, slice: &mut [f16], kind: SoftmaxKind) -> TractResult<()> {
        let max = Func::ReduceMax.reduce_f16()?.run(slice)?;
        match kind {
            SoftmaxKind::Softmax => {
                // f16 softmax runs the scalar libm exp and accumulates the sum
                // in f32: a row long enough for the running sum to outgrow its
                // own terms stalls in f16, costing double digits of relative
                // error by a few thousand elements. A vectorized f16 exp is not
                // worth its fragility, f16 softmax always lands here.
                let mut sum = 0f32;
                slice.iter_mut().for_each(|x| {
                    *x = (*x - max).exp();
                    sum += x.to_f32();
                });
                let rsum = f16::from_f32(sum.recip());
                Func::MulByScalar.ew_f16_param()?.run_with_params(slice, rsum)?;
            }
            SoftmaxKind::LogSoftmax => {
                let mut exp_sum = 0f32;
                slice.iter_mut().for_each(|x| {
                    *x -= max;
                    exp_sum += x.exp().to_f32();
                });
                let log_sum = f16::from_f32(exp_sum.ln());
                slice.iter_mut().for_each(|x| *x -= log_sum);
            }
        }
        Ok(())
    }
}

/// The f32 row kernels, built once per evaluation: a softmax over many short rows would
/// otherwise build them again for every row.
struct F32Kernels {
    max: Box<dyn tract_linalg::reduce::Reduce<f32>>,
    exp_sum: Box<dyn tract_linalg::reduce::MapReduce<f32, f32>>,
    scale: Box<dyn tract_linalg::element_wise::ElementWise<f32, f32>>,
}

impl F32Kernels {
    fn new() -> TractResult<F32Kernels> {
        Ok(F32Kernels {
            max: Func::ReduceMax.reduce_f32()?,
            exp_sum: Func::Softmax2.map_reduce_f32()?,
            scale: Func::MulByScalar.ew_f32_param()?,
        })
    }
}

fn softmax_inner_slice_f32(
    slice: &mut [f32],
    kind: SoftmaxKind,
    kernels: &F32Kernels,
) -> TractResult<()> {
    let max = kernels.max.run(slice)?;
    match kind {
        SoftmaxKind::Softmax => {
            let sum = kernels.exp_sum.run_with_params(slice, max)?;
            kernels.scale.run_with_params(slice, sum.recip())?;
        }
        SoftmaxKind::LogSoftmax => {
            let mut exp_sum = f32::zero();
            slice.iter_mut().for_each(|x| {
                *x -= max;
                exp_sum += x.exp();
            });
            let log_sum = exp_sum.ln();
            slice.iter_mut().for_each(|x| *x -= log_sum);
        }
    }
    Ok(())
}

fn softmax_inner<T: Float + Datum + std::iter::Sum, D: Dimension>(
    mut view: ArrayViewMut<T, D>,
    kind: SoftmaxKind,
) {
    let max =
        *view.iter().max_by(|i, j| i.partial_cmp(j).unwrap_or(std::cmp::Ordering::Less)).unwrap();
    view.mapv_inplace(|x| x - max);
    let exp_sum = view.iter().map(|&x| x.exp()).sum();
    match kind {
        SoftmaxKind::Softmax => {
            view.mapv_inplace(|x| x.exp() / exp_sum);
        }
        SoftmaxKind::LogSoftmax => {
            let log_sum = exp_sum.ln();
            view.mapv_inplace(|x| x - log_sum);
        }
    }
}

fn softmax_quant_inner<D: Dimension>(
    mut view: ArrayViewMut<u8, D>,
    src_is_signed: bool,
    in_qp: QParams,
    out_is_signed: bool,
    out_qp: QParams,
) {
    let (_, in_scale) = in_qp.zp_scale();
    let (scale_in_multiplier, scale_in_shift) = convert_scale_to_mult_shift(in_scale).unwrap();
    let (_, out_scale) = out_qp.zp_scale();
    let (scale_out_multiplier, scale_out_shift) = convert_scale_to_mult_shift(out_scale).unwrap();
    let shift = 26 - scale_in_shift;

    // Compute the exponentials x - max
    let mut buffer = vec![0_i32; view.len()];

    // Handle the case were we considered an i8 as an u8 and still get the right x - max.
    let safe_u8 = if src_is_signed { |x: &u8| x.wrapping_add(128) } else { |x: &u8| *x };

    let max = view.iter().map(safe_u8).max().unwrap();
    view.iter().zip(buffer.iter_mut()).for_each(|(x, exp)| {
        let input_diff = safe_u8(x) as i32 - max as i32;

        // We scale the input to be in Q5_26
        let scaled_input_diff = if scale_in_multiplier != 0 {
            saturating_rounding_multiply_by_pot(
                saturating_rounding_doubling_high_mul(input_diff, scale_in_multiplier),
                shift as i32,
            )
        } else {
            saturating_rounding_multiply_by_pot(input_diff, shift as i32)
        };

        // It expects an input from Q5_26 and returns an output in Q0_31
        *exp = exp_on_negative_values(scaled_input_diff);
    });

    // Compute sum of exp
    // The sum is stored as an Q12_19 that's why we need to recale from Q0_31 to Q12_19 before summing.
    let sum_of_exp = buffer.iter().map(|it| rescale(*it, 0, 12)).sum();

    // Compute 1/sum_of_exp
    // The result of this function is in Q0_31
    let (inv_sum_of_exp, num_bits_over_unit) = get_reciprocal(sum_of_exp, 12);

    // Compute the exponent value needed to be in Q24_8 before the final rescaling
    let exponent = num_bits_over_unit as isize + 31 - 8;

    view.iter_mut().zip(buffer.iter()).for_each(|(it, exp)| {
        // Compute the product of exp * 1/sum_of_exp and scale the result in Q24_8
        let unsat_output = rounding_divide_by_pot(
            saturating_rounding_doubling_high_mul(inv_sum_of_exp, *exp),
            exponent as i32,
        );

        // Scale the final result in the output scale range
        let unsat_scaled_output = {
            if scale_out_multiplier != 0 {
                let (inv_multiplier, num_bits) = get_reciprocal(scale_out_multiplier, 1);
                rounding_divide_by_pot(
                    saturating_rounding_doubling_high_mul(unsat_output, inv_multiplier),
                    (8 - scale_out_shift - 1 - num_bits as isize) as i32,
                )
            } else {
                rounding_divide_by_pot(unsat_output, (8 - scale_out_shift) as i32)
            }
        };

        // Return the final result by clipping the computed value within its range
        // and casting it to u8 in any case.
        #[allow(unknown_lints, unnecessary_transmutes)]
        if out_is_signed {
            *it = unsafe {
                std::mem::transmute::<i8, u8>(i32::max(
                    i32::min(unsat_scaled_output, i8::MAX as i32),
                    i8::MIN as i32,
                ) as i8)
            };
        } else {
            *it = i32::max(i32::min(unsat_scaled_output, u8::MAX as i32), u8::MIN as i32) as u8;
        }
    });
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::ops::nn::DataFormat::NCHW;
    use anyhow::Result;
    use num_traits::PrimInt;
    use proptest::collection::vec;
    use proptest::prelude::*;
    use tract_data::internal::QParams::ZpScale;

    fn assert_is_close(found: f32, expected: f32, in_dt: DatumType, out_dt: DatumType) {
        let (_, in_epsilon) = in_dt.zp_scale();
        let (_, out_epsilon) = out_dt.zp_scale();
        let epsilon = in_epsilon + out_epsilon;
        let error = (found - expected).abs();
        assert!(
            error <= epsilon,
            "epsilon eq failed: |{found:?}-{expected:?}|={error} should be <= {epsilon}"
        );
    }

    // Generate a random tensor with a quantized datum type
    fn qtensor<T: PrimInt + Datum + Arbitrary>(shape: Vec<usize>) -> BoxedStrategy<Tensor> {
        let len = shape.iter().product::<usize>();
        let dt = q_datum::<T>((0.0001f32..0.1).boxed());
        (vec(any::<T>(), len..=len), dt)
            .prop_map(move |(vec, dt)| (ArrayD::from_shape_vec(shape.clone(), vec).unwrap(), dt))
            .prop_map(move |(array, dt)| {
                let mut tensor = array.into_tensor();
                unsafe { tensor.set_datum_type(dt) };
                tensor
            })
            .boxed()
    }

    // Generate a random quantized datum type
    fn q_datum<T: PrimInt + Datum>(range: BoxedStrategy<f32>) -> BoxedStrategy<DatumType> {
        let max_integer_bits = std::mem::size_of::<T>() * 8 - T::datum_type().is_signed() as usize;
        prop_oneof![
            (1usize..max_integer_bits).prop_map(|fixed_point| { 2f32.powi(-(fixed_point as i32)) }),
            range
        ]
        .prop_map(|scale| {
            if T::datum_type().is_signed() {
                DatumType::QI8(ZpScale { zero_point: 0, scale })
            } else {
                DatumType::QU8(ZpScale { zero_point: 0, scale })
            }
        })
        .boxed()
    }

    #[derive(Debug)]
    struct SoftmaxProblem {
        data: Tensor,
        axes: TVec<usize>,
        output_dt: DatumType,
    }

    impl SoftmaxProblem {
        fn check(&self) -> Result<()> {
            let inputs = tvec!(self.data.clone().into_tvalue());
            let quant_output_dt = Some(self.output_dt).filter(|dt| !dt.is_float());
            let softmax =
                Softmax { axes: self.axes.clone(), quant_output_dt, ..Softmax::default() };

            // Compute quantized output
            let result = softmax.eval(&EvalContext::out_of_plan(), inputs)?;
            let result = args_1!(result);
            let result_float = result.cast_to::<f32>()?;

            // Compute reference output
            let input_float = self.data.cast_to::<f32>()?;
            let inputs_float = tvec!(input_float.into_owned().into_tvalue());
            let softmax_float = Softmax { axes: self.axes.clone(), ..Softmax::default() };
            let reference_float = softmax_float.eval(&EvalContext::out_of_plan(), inputs_float)?;
            let reference_array = args_1!(reference_float);
            let reference = reference_array.to_plain_array_view::<f32>()?;

            result_float
                .to_plain_array_view::<f32>()?
                .iter()
                .zip(reference.iter())
                .for_each(|(a, b)| assert_is_close(*a, *b, self.data.datum_type(), self.output_dt));
            Ok(())
        }
    }

    impl Arbitrary for SoftmaxProblem {
        type Parameters = ();
        type Strategy = BoxedStrategy<SoftmaxProblem>;
        fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
            (1usize..2, 1usize..2, 1usize..5, 1usize..5, 0usize..4)
                .prop_flat_map(|(n, c, h, w, axis)| {
                    let shape_in: Vec<usize> =
                        NCHW.from_n_c_hw(n, c, [h, w]).unwrap().shape.to_vec();
                    (
                        prop_oneof![qtensor::<i8>(shape_in.clone()), qtensor::<u8>(shape_in)],
                        Just(tvec![axis]),
                        prop_oneof![
                            q_datum::<u8>((0.008f32..0.1).boxed()),
                            q_datum::<i8>((0.008f32..0.1).boxed())
                        ],
                    )
                })
                .prop_map(|(data, axes, output_dt)| SoftmaxProblem { data, axes, output_dt })
                .boxed()
        }
    }

    #[derive(Debug)]
    pub struct InnerSoftmaxProblem {
        in_qp: QParams,
        out_qp: QParams,
        data: Vec<i8>,
    }

    impl InnerSoftmaxProblem {
        fn check(&self) -> Result<()> {
            let quantized = self.quantized();
            let reference = self.reference();
            assert!(quantized.iter().zip(reference.iter()).all(|(quantized, expected)| {
                let abs_diff = if *quantized > *expected {
                    quantized - *expected
                } else {
                    expected - *quantized
                };
                abs_diff <= 1
            }));
            Ok(())
        }

        fn reference(&self) -> Vec<u8> {
            let (in_zero_point, in_scale) = self.in_qp.zp_scale();
            let (out_zero_point, out_scale) = self.out_qp.zp_scale();
            let in_float =
                self.data.iter().map(|it| (*it as f32 - in_zero_point as f32) * in_scale).collect();
            let mut in_float_array = Array1::from_vec(in_float);
            softmax_inner(in_float_array.view_mut(), SoftmaxKind::default());
            in_float_array
                .iter()
                .map(|it| {
                    ((*it / out_scale).round() as i32 + out_zero_point)
                        .max(u8::MIN as i32)
                        .min(u8::MAX as i32) as u8
                })
                .collect()
        }

        fn quantized(&self) -> Vec<u8> {
            let in_data: Vec<u8> = unsafe { std::mem::transmute(self.data.clone()) };
            let mut in_array = Array1::from_vec(in_data);
            softmax_quant_inner(in_array.view_mut(), true, self.in_qp, false, self.out_qp);
            in_array.to_vec()
        }
    }

    impl Arbitrary for InnerSoftmaxProblem {
        type Parameters = ();
        type Strategy = BoxedStrategy<InnerSoftmaxProblem>;
        fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
            (
                prop_oneof![
                    q_datum::<i8>((0.0001f32..0.01).boxed()),
                    q_datum::<u8>((0.0001f32..0.01).boxed())
                ],
                prop_oneof![
                    q_datum::<u8>((0.008f32..0.1).boxed()),
                    q_datum::<i8>((0.008f32..0.1).boxed())
                ],
                vec(any::<i8>(), 1..10),
            )
                .prop_map(|(in_qp, out_qp, data)| InnerSoftmaxProblem {
                    in_qp: in_qp.qparams().unwrap(),
                    out_qp: out_qp.qparams().unwrap(),
                    data,
                })
                .boxed()
        }
    }

    proptest::proptest! {
        #![proptest_config(ProptestConfig::with_cases(1000))]
        #[test]
        fn test_softmax_inner_prop(pb in any::<InnerSoftmaxProblem>()) {
            pb.check().unwrap()
        }
    }

    proptest::proptest! {
        #![proptest_config(ProptestConfig::with_cases(1000))]
        #[test]
        fn test_softmax_prop(pb in any::<SoftmaxProblem>()) {
            pb.check().unwrap()
        }
    }

    #[test]
    // We test QU8 -> QU8
    fn test_softmax_trivial_0() -> Result<()> {
        let input_dt = DatumType::QU8(ZpScale { zero_point: 0, scale: 0.03125 }); // Q3_5
        let output_dt = DatumType::QU8(ZpScale { zero_point: 0, scale: 0.00390625 }); // Q0_8;
        let mut data = Tensor::from_shape(&[1, 1, 2, 2], &[0_u8, 0, 0, 4])?;
        unsafe { data.set_datum_type(input_dt) };

        let prob = SoftmaxProblem { data, axes: tvec![3], output_dt };
        prob.check()?;
        Ok(())
    }

    #[test]
    // We test QI8 -> QU8
    fn test_softmax_trivial_1() -> Result<()> {
        let input_dt = DatumType::QI8(ZpScale { zero_point: 0, scale: 0.0625 }); // Q3_4
        let output_dt = DatumType::QU8(ZpScale { zero_point: 0, scale: 0.00390625 }); // Q0_8;
        let mut data = Tensor::from_shape(&[1, 1, 2, 2], &[0_i8, 0, 0, 4])?;
        unsafe { data.set_datum_type(input_dt) };

        let prob = SoftmaxProblem { data, axes: tvec![3], output_dt };
        prob.check()?;
        Ok(())
    }

    #[test]
    // We test QI8 -> QI8
    fn test_softmax_trivial_2() -> Result<()> {
        let input_dt = DatumType::QI8(ZpScale { zero_point: 0, scale: 0.0625 }); // Q3_4
        let output_dt = DatumType::QI8(ZpScale { zero_point: 0, scale: 0.0078125 }); // Q0_7;
        let mut data = Tensor::from_shape(&[1, 1, 2, 2], &[0_i8, 0, 0, -4])?;
        unsafe { data.set_datum_type(input_dt) };

        let prob = SoftmaxProblem { data, axes: tvec![3], output_dt };
        prob.check()?;
        Ok(())
    }

    #[test]
    // We test QU8 -> QI8
    fn test_softmax_trivial_3() -> Result<()> {
        let input_dt = DatumType::QU8(ZpScale { zero_point: 0, scale: 0.03125 }); // Q3_5
        let output_dt = DatumType::QI8(ZpScale { zero_point: 0, scale: 0.0078125 }); // Q0_7;
        let mut data = Tensor::from_shape(&[1, 1, 2, 2], &[0_u8, 0, 0, 4])?;
        unsafe { data.set_datum_type(input_dt) };

        let prob = SoftmaxProblem { data, axes: tvec![2], output_dt };
        prob.check()?;
        Ok(())
    }

    #[test]
    fn test_softmax_1() -> Result<()> {
        let input_dt = DatumType::QI8(ZpScale { zero_point: 0, scale: 0.5 }); // Q6_1
        let output_dt = DatumType::QU8(ZpScale { zero_point: 0, scale: 0.5 }); // Q7_1
        let mut data = Tensor::from_shape(&[1, 1, 1, 2], &[115_i8, 115])?;
        unsafe { data.set_datum_type(input_dt) };

        let prob = SoftmaxProblem { data, axes: tvec![3], output_dt };
        prob.check()?;
        Ok(())
    }

    #[test]
    fn test_softmax_2() -> Result<()> {
        let input_dt = DatumType::QI8(ZpScale { zero_point: 0, scale: 0.0001 });
        let output_dt = DatumType::QU8(ZpScale { zero_point: 0, scale: 0.008 });
        let mut data = Tensor::from_shape(&[1, 1, 1, 2], &[115_i8, 115])?;
        unsafe { data.set_datum_type(input_dt) };

        let prob = SoftmaxProblem { data, axes: tvec![3], output_dt };
        prob.check()?;
        Ok(())
    }

    #[test]
    fn test_softmax_3() -> Result<()> {
        let input_dt = DatumType::QU8(ZpScale { zero_point: 0, scale: 0.6220956 });
        let output_dt = DatumType::QU8(ZpScale { zero_point: 0, scale: 0.5187921 });
        let mut data = Tensor::from_shape(&[1, 1, 1, 2], &[13_u8, 218])?;
        unsafe { data.set_datum_type(input_dt) };

        let prob = SoftmaxProblem { data, axes: tvec![3], output_dt };
        prob.check()?;
        Ok(())
    }

    #[test]
    fn test_inner_softmax_1() -> Result<()> {
        let in_qp = ZpScale { zero_point: 0, scale: 0.03125 };
        let out_qp = ZpScale { zero_point: 0, scale: 0.5 };
        let data = vec![0_i8, 1];

        let prob = InnerSoftmaxProblem { in_qp, out_qp, data };
        prob.check()?;
        Ok(())
    }

    #[test]
    fn test_inner_softmax_2() -> Result<()> {
        let in_qp = ZpScale { zero_point: 0, scale: 0.5 };
        let out_qp = ZpScale { zero_point: 0, scale: 0.03125 };
        let data = vec![100i8, -28];

        let prob = InnerSoftmaxProblem { in_qp, out_qp, data };
        prob.check()?;
        Ok(())
    }

    #[test]
    fn test_inner_softmax_not_pow_2_1() -> Result<()> {
        let in_qp = ZpScale { zero_point: 0, scale: 0.7298456 };
        let out_qp = ZpScale { zero_point: 0, scale: 0.03125 };
        let data = vec![100i8, -28];

        let prob = InnerSoftmaxProblem { in_qp, out_qp, data };
        prob.check()?;
        Ok(())
    }

    #[test]
    #[ignore]
    // Fails but the difference is quite low and the sum still give exactly one:
    // quantized: 110(0.88), 15(0.12)
    // expected: 112(0.896), 13(0.104)
    fn test_inner_softmax_not_pow_2_2() -> Result<()> {
        let in_qp = ZpScale { zero_point: 0, scale: 0.2123116 };
        let out_qp = ZpScale { zero_point: 0, scale: 0.008 };
        let data = vec![118i8, 108];

        let prob = InnerSoftmaxProblem { in_qp, out_qp, data };
        prob.check()?;
        Ok(())
    }

    #[test]
    #[ignore]
    // Fails but the difference is quite low and the sum still give exactly one:
    // quantized: 40(0.625), 24(0.375)
    // expected: 42(0.65625), 22(0.34375)
    fn test_inner_softmax_not_pow_2_3() -> Result<()> {
        let in_qp = ZpScale { zero_point: 0, scale: 0.33034274 };
        let out_qp = ZpScale { zero_point: 0, scale: 0.015625 };
        let data = vec![45i8, 43];

        let prob = InnerSoftmaxProblem { in_qp, out_qp, data };
        prob.check()?;
        Ok(())
    }
}

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

    /// The f16 row must stay close to the same softmax evaluated in f32. A row
    /// long enough for the running sum to outgrow its own terms is the case an
    /// f16 accumulator silently drops.
    #[test]
    fn long_rows_keep_their_normalisation() -> TractResult<()> {
        for len in [1024usize, 4096, 8192] {
            let logits: Vec<f32> = (0..len).map(|i| ((i as f32) * 0.7).sin() * 0.5).collect();

            let mut half: Vec<f16> = logits.iter().map(|v| f16::from_f32(*v)).collect();
            Softmax::new(tvec![0], None, SoftmaxKind::Softmax)
                .softmax_inner_slice_f16(&mut half, SoftmaxKind::Softmax)?;

            let max = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
            let exps: Vec<f32> = logits.iter().map(|v| (v - max).exp()).collect();
            let total: f32 = exps.iter().sum();

            let got: f32 = half.iter().map(|v| v.to_f32()).sum();
            assert!((got - 1.0).abs() < 0.02, "len {len}: row sums to {got}, expected 1.0");
            for (i, (h, e)) in half.iter().zip(&exps).enumerate() {
                let want = e / total;
                let err = (h.to_f32() - want).abs() / want.max(f32::MIN_POSITIVE);
                assert!(err < 0.05, "len {len} lane {i}: got {h}, want {want}, rel {err}");
            }
        }
        Ok(())
    }
}