rustyml 0.16.0

//! Element-wise activation functions and the standalone activation layers
//!
//! Defines the [`Activation`] enum, the single source of truth for each activation's forward
//! transform and derivative. The thin layer wrappers delegate their math to it: [`Linear`],
//! [`ReLU`], [`LeakyReLU`], [`ELU`], [`SELU`], [`Sigmoid`], [`HardSigmoid`], [`Tanh`],
//! [`Softplus`], [`Softsign`], [`Exponential`], and [`Softmax`]
//!
//! [`PReLU`] stands apart. Its negative-side slope is a trainable array, so it holds its own
//! math and its own weights instead of delegating to the enum

use crate::error::{Context, Error};
use crate::neural_network::Tensor;
use crate::parallel_gates::{cheap_map_parallel_threshold, exp_map_parallel_threshold};
use crate::{Deserialize, Serialize};
use ndarray::{Array2, ArrayView1, ArrayViewMut1, Axis, Zip};
use rayon::iter::{IntoParallelIterator, ParallelIterator};

/// ELU (Exponential Linear Unit) activation layer
pub mod elu;
/// Exponential activation layer
pub mod exponential;
/// Hard sigmoid activation layer, a piecewise-linear approximation of the logistic sigmoid
pub mod hard_sigmoid;
/// Leaky ReLU activation layer
pub mod leaky_relu;
/// Linear (Identity) activation layer
pub mod linear;
/// PReLU activation layer, whose negative-side slope is a trainable parameter
pub mod p_relu;
/// ReLU (Rectified Linear Unit) activation layer
pub mod relu;
/// SELU (Scaled Exponential Linear Unit) activation layer
pub mod selu;
/// Sigmoid activation layer
pub mod sigmoid;
/// Softmax activation layer
pub mod softmax;
/// Softplus activation layer
pub mod softplus;
/// Softsign activation layer
pub mod softsign;
/// Tanh (Hyperbolic Tangent) activation layer
pub mod tanh;

pub use elu::ELU;
pub use exponential::Exponential;
pub use hard_sigmoid::HardSigmoid;
pub use leaky_relu::LeakyReLU;
pub use linear::Linear;
pub use p_relu::PReLU;
pub use relu::ReLU;
pub use selu::SELU;
pub use sigmoid::Sigmoid;
pub use softmax::Softmax;
pub use softplus::Softplus;
pub use softsign::Softsign;
pub use tanh::Tanh;

/// SELU's fixed `alpha`, from Klambauer et al. (2017)
///
/// The published constant is 1.6732632423543772848170429916717. This literal is the f32 it
/// rounds to
const SELU_ALPHA: f32 = 1.673_263_2;

/// SELU's fixed `scale`, from Klambauer et al. (2017)
///
/// The published constant is 1.0507009873554804934193349852946. This literal is the f32 it
/// rounds to
const SELU_SCALE: f32 = 1.050_701;

/// The product `SELU_SCALE * SELU_ALPHA`, about 1.7580993
///
/// The negative branch multiplies by this single constant. The backward pass then recovers
/// the derivative as `a + SELU_SCALE_ALPHA`, which is exact for the value the forward pass
/// wrote
const SELU_SCALE_ALPHA: f32 = SELU_SCALE * SELU_ALPHA;

/// The slope of the hard sigmoid's linear segment, `1/6`
const HARD_SIGMOID_SLOPE: f32 = 1.0 / 6.0;

/// The default softmax axis, the last axis of the input
///
/// A negative axis counts back from the end, so `-1` holds for every rank
pub const DEFAULT_SOFTMAX_AXIS: i32 = -1;

/// The element-wise activation functions that trainable layers can embed
///
/// Dense, the convolutional layers, and the recurrent layers each carry an `Activation` value
/// instead of a generic activation type parameter. A runtime enum keeps the host layers
/// non-generic. This removes monomorphization bloat and lets weight deserialization downcast
/// every layer to a single concrete type, instead of probing each `Layer<Act>` pairing
///
/// Every standalone activation *layer* in this module delegates its math here. This enum is
/// the single source of truth for both the forward transform and its derivative. The
/// algorithm lives here, and not in a layer `impl` block
///
/// # Notes
///
/// [`Activation::backward`] receives the *activated output* `a`, never the pre-activation `z`.
/// Every variant's derivative is therefore expressed through `a` alone, and a host layer caches
/// 1 tensor rather than 2. This is the output-only derivative contract that every variant below
/// follows.
///
/// The contract admits every activation whose derivative has a closed form in `a`, which covers
/// the full family below. It excludes GELU, SiLU (Swish), and Mish, whose `a = z * g(z)` shape
/// has no closed-form inverse. Supporting them needs a wider contract that also hands the
/// backward pass the pre-activation.
///
/// The 2 saturating variants pay a small accuracy cost for the contract. `ELU` recovers its
/// negative branch as `a + alpha`, and `SELU` recovers its negative branch as
/// `a + SELU_SCALE_ALPHA`. Both computations subtract 2 near-equal values far down the tail.
/// The absolute error stays inside 1 unit in the last place of that constant, at inputs where
/// the derivative is already close to 0.
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub enum Activation {
    /// Identity activation, `f(x) = x`
    Linear,
    /// Rectified linear unit, `max(0, x)`
    ReLU,
    /// Logistic sigmoid, `1 / (1 + e^-x)`
    Sigmoid,
    /// Hyperbolic tangent, `tanh(x)`
    Tanh,
    /// Softmax over 1 axis, `e^x / sum(e^x)` along that axis
    ///
    /// The lanes along the chosen axis each become a probability distribution that sums
    /// to 1. The whole tensor keeps its shape
    Softmax {
        /// Axis to normalize. A negative value counts back from the end, so the default
        /// `-1` always means the last axis
        ///
        /// The axis resolves against the rank of the input on each call, not when the
        /// value is built. The same value therefore normalizes the last axis of a rank-2
        /// input and the last axis of a rank-4 input. A resolved index outside the rank
        /// of the input fails the call
        ///
        /// A trainable host layer accepts `-1` alone. See [`Activation::validate`]
        axis: i32,
    },
    /// Leaky rectified linear unit, `x` for `x >= 0` and `negative_slope * x` below it
    ///
    /// Unlike [`Activation::ReLU`], the negative side keeps a non-zero gradient, so a unit
    /// whose pre-activation stays negative can still recover
    LeakyReLU {
        /// Slope applied below 0. Must be finite and greater than 0. The layer form
        /// defaults to `0.3`. Use [`Activation::ReLU`] for a slope of 0
        negative_slope: f32,
    },
    /// Exponential linear unit, `x` for `x > 0` and `alpha * (e^x - 1)` below it
    ELU {
        /// Scale of the saturating negative branch. Must be finite and greater than 0.
        /// The layer form defaults to `1.0`
        alpha: f32,
    },
    /// Scaled exponential linear unit, `scale * x` for `x > 0` and
    /// `scale * alpha * (e^x - 1)` below it
    ///
    /// `alpha` and `scale` are the fixed constants of Klambauer et al. (2017). Pair it with
    /// Lecun-normal initialization for the self-normalizing property to hold
    SELU,
    /// Softplus, `ln(1 + e^x)`, a smooth approximation of [`Activation::ReLU`]
    Softplus,
    /// Softsign, `x / (1 + |x|)`, a bounded activation that saturates polynomially rather
    /// than exponentially
    Softsign,
    /// Hard sigmoid, `clip(x/6 + 0.5, 0, 1)`, a piecewise-linear approximation of
    /// [`Activation::Sigmoid`] with no exponential
    HardSigmoid,
    /// Exponential, `e^x`
    Exponential,
}

impl Activation {
    /// Applies the activation to a pre-activation tensor `z` and returns the activated output
    ///
    /// # Parameters
    ///
    /// - `z` - Pre-activation tensor (the linear output of the host layer)
    ///
    /// # Returns
    ///
    /// - `Result<Tensor, Error>` - Activated tensor with the same shape as `z`
    ///
    /// # Errors
    ///
    /// - `Error::InvalidInput` - The Softmax axis resolves outside the rank of `z`
    /// - `Error::Computation` - Softmax failed to reshape the input
    pub fn forward(&self, z: &Tensor) -> Result<Tensor, Error> {
        match self {
            Activation::Linear => Ok(z.clone()),
            Activation::ReLU => {
                let relu = |x: f32| if x <= 0.0 { 0.0 } else { x };
                let out = if z.len() >= cheap_map_parallel_threshold() {
                    Zip::from(z).par_map_collect(|&x| relu(x))
                } else {
                    z.mapv(relu)
                };
                Ok(out)
            }
            Activation::Sigmoid => {
                let sigmoid = |x: f32| 1.0 / (1.0 + (-x).exp());
                let out = if z.len() >= exp_map_parallel_threshold() {
                    Zip::from(z).par_map_collect(|&x| sigmoid(x))
                } else {
                    z.mapv(sigmoid)
                };
                Ok(out)
            }
            Activation::Tanh => {
                let tanh = |x: f32| x.tanh();
                let out = if z.len() >= exp_map_parallel_threshold() {
                    Zip::from(z).par_map_collect(|&x| tanh(x))
                } else {
                    z.mapv(tanh)
                };
                Ok(out)
            }
            Activation::Softmax { axis } => softmax_forward(z, *axis),
            Activation::LeakyReLU { negative_slope } => {
                let slope = *negative_slope;
                let leaky_relu = |x: f32| if x >= 0.0 { x } else { slope * x };
                let out = if z.len() >= cheap_map_parallel_threshold() {
                    Zip::from(z).par_map_collect(|&x| leaky_relu(x))
                } else {
                    z.mapv(leaky_relu)
                };
                Ok(out)
            }
            Activation::ELU { alpha } => {
                let alpha = *alpha;
                // `exp_m1` keeps the full precision of `e^x - 1` as x approaches 0
                let elu = |x: f32| if x > 0.0 { x } else { alpha * x.exp_m1() };
                let out = if z.len() >= exp_map_parallel_threshold() {
                    Zip::from(z).par_map_collect(|&x| elu(x))
                } else {
                    z.mapv(elu)
                };
                Ok(out)
            }
            Activation::SELU => {
                let selu = |x: f32| {
                    if x > 0.0 {
                        SELU_SCALE * x
                    } else {
                        SELU_SCALE_ALPHA * x.exp_m1()
                    }
                };
                let out = if z.len() >= exp_map_parallel_threshold() {
                    Zip::from(z).par_map_collect(|&x| selu(x))
                } else {
                    z.mapv(selu)
                };
                Ok(out)
            }
            Activation::Softplus => {
                // Factoring `e^x` out of the logarithm above 0 keeps both tails in range.
                // The direct form overflows for large x, and loses every digit for small x
                let softplus = |x: f32| {
                    if x > 0.0 {
                        x + (-x).exp().ln_1p()
                    } else {
                        x.exp().ln_1p()
                    }
                };
                let out = if z.len() >= exp_map_parallel_threshold() {
                    Zip::from(z).par_map_collect(|&x| softplus(x))
                } else {
                    z.mapv(softplus)
                };
                Ok(out)
            }
            Activation::Softsign => {
                let softsign = |x: f32| x / (1.0 + x.abs());
                let out = if z.len() >= cheap_map_parallel_threshold() {
                    Zip::from(z).par_map_collect(|&x| softsign(x))
                } else {
                    z.mapv(softsign)
                };
                Ok(out)
            }
            Activation::HardSigmoid => {
                let hard_sigmoid = |x: f32| (x + 3.0).clamp(0.0, 6.0) / 6.0;
                let out = if z.len() >= cheap_map_parallel_threshold() {
                    Zip::from(z).par_map_collect(|&x| hard_sigmoid(x))
                } else {
                    z.mapv(hard_sigmoid)
                };
                Ok(out)
            }
            Activation::Exponential => {
                let exponential = |x: f32| x.exp();
                let out = if z.len() >= exp_map_parallel_threshold() {
                    Zip::from(z).par_map_collect(|&x| exponential(x))
                } else {
                    z.mapv(exponential)
                };
                Ok(out)
            }
        }
    }

    /// Computes the gradient with respect to the pre-activation input
    ///
    /// Every supported activation's derivative is expressible in terms of its own
    /// output, so this takes the cached activated tensor rather than the original input
    ///
    /// This is pure math with no clamping or NaN/Inf sanitization
    ///
    /// # Parameters
    ///
    /// - `activated` - The activated output `a` produced by [`forward`](Activation::forward)
    /// - `grad_output` - Upstream gradient `dL/da`
    ///
    /// # Returns
    ///
    /// - `Result<Tensor, Error>` - The gradient `dL/dz`, same shape as `activated`
    ///
    /// # Errors
    ///
    /// - `Error::InvalidInput` - The Softmax axis resolves outside the rank of `activated`
    /// - `Error::ShapeMismatch` - Softmax received a `grad_output` whose shape differs from
    ///   the shape of `activated`
    /// - `Error::Computation` - Softmax failed to reshape the tensors
    pub fn backward(&self, activated: &Tensor, grad_output: &Tensor) -> Result<Tensor, Error> {
        match self {
            Activation::Linear => Ok(grad_output.clone()),
            Activation::ReLU => {
                // ReLU'(z) = 1 when z > 0. Since a = max(0, z), `a > 0` exactly when `z > 0`
                let mut grad = grad_output.clone();
                let relu_grad = |g: &mut f32, &a: &f32| {
                    if a <= 0.0 {
                        *g = 0.0;
                    }
                };
                if activated.len() >= cheap_map_parallel_threshold() {
                    Zip::from(&mut grad).and(activated).par_for_each(relu_grad);
                } else {
                    Zip::from(&mut grad).and(activated).for_each(relu_grad);
                }
                Ok(grad)
            }
            Activation::Sigmoid => {
                let mut grad = grad_output.clone();
                let sigmoid_grad = |g: &mut f32, &a: &f32| {
                    *g *= a * (1.0 - a);
                };
                if grad.len() >= cheap_map_parallel_threshold() {
                    Zip::from(&mut grad)
                        .and(activated)
                        .par_for_each(sigmoid_grad);
                } else {
                    Zip::from(&mut grad).and(activated).for_each(sigmoid_grad);
                }
                Ok(grad)
            }
            Activation::Tanh => {
                let mut grad = grad_output.clone();
                let tanh_grad = |g: &mut f32, &a: &f32| {
                    *g *= 1.0 - a * a;
                };
                if activated.len() >= cheap_map_parallel_threshold() {
                    Zip::from(&mut grad).and(activated).par_for_each(tanh_grad);
                } else {
                    Zip::from(&mut grad).and(activated).for_each(tanh_grad);
                }
                Ok(grad)
            }
            Activation::Softmax { axis } => softmax_backward(activated, grad_output, *axis),
            Activation::LeakyReLU { negative_slope } => {
                // LeakyReLU'(z) = 1 when z >= 0, and `negative_slope` below it. A positive
                // slope keeps the sign of z, so `a < 0` marks exactly the elements with z < 0
                let slope = *negative_slope;
                let mut grad = grad_output.clone();
                let leaky_relu_grad = |g: &mut f32, &a: &f32| {
                    if a < 0.0 {
                        *g *= slope;
                    }
                };
                if activated.len() >= cheap_map_parallel_threshold() {
                    Zip::from(&mut grad)
                        .and(activated)
                        .par_for_each(leaky_relu_grad);
                } else {
                    Zip::from(&mut grad)
                        .and(activated)
                        .for_each(leaky_relu_grad);
                }
                Ok(grad)
            }
            Activation::ELU { alpha } => {
                // ELU'(z) = 1 when z > 0. Below that a = alpha * (e^z - 1), which rearranges
                // to alpha * e^z = a + alpha, so the derivative needs no exponential
                let alpha = *alpha;
                let mut grad = grad_output.clone();
                let elu_grad = |g: &mut f32, &a: &f32| {
                    if a <= 0.0 {
                        *g *= a + alpha;
                    }
                };
                if activated.len() >= cheap_map_parallel_threshold() {
                    Zip::from(&mut grad).and(activated).par_for_each(elu_grad);
                } else {
                    Zip::from(&mut grad).and(activated).for_each(elu_grad);
                }
                Ok(grad)
            }
            Activation::SELU => {
                // SELU'(z) = scale when z > 0, and scale * alpha * e^z below it. The same
                // rearrangement as ELU turns that into a + SELU_SCALE_ALPHA
                let mut grad = grad_output.clone();
                let selu_grad = |g: &mut f32, &a: &f32| {
                    *g *= if a > 0.0 {
                        SELU_SCALE
                    } else {
                        a + SELU_SCALE_ALPHA
                    };
                };
                if activated.len() >= cheap_map_parallel_threshold() {
                    Zip::from(&mut grad).and(activated).par_for_each(selu_grad);
                } else {
                    Zip::from(&mut grad).and(activated).for_each(selu_grad);
                }
                Ok(grad)
            }
            Activation::Softplus => {
                // softplus'(z) = sigmoid(z). With a = ln(1 + e^z), that is 1 - e^-a.
                // `exp_m1` holds the tiny values the far negative tail produces
                let mut grad = grad_output.clone();
                let softplus_grad = |g: &mut f32, &a: &f32| {
                    *g *= -(-a).exp_m1();
                };
                if activated.len() >= exp_map_parallel_threshold() {
                    Zip::from(&mut grad)
                        .and(activated)
                        .par_for_each(softplus_grad);
                } else {
                    Zip::from(&mut grad).and(activated).for_each(softplus_grad);
                }
                Ok(grad)
            }
            Activation::Softsign => {
                // softsign'(z) = 1 / (1 + |z|)^2. With a = z / (1 + |z|), that is (1 - |a|)^2
                let mut grad = grad_output.clone();
                let softsign_grad = |g: &mut f32, &a: &f32| {
                    let t = 1.0 - a.abs();
                    *g *= t * t;
                };
                if activated.len() >= cheap_map_parallel_threshold() {
                    Zip::from(&mut grad)
                        .and(activated)
                        .par_for_each(softsign_grad);
                } else {
                    Zip::from(&mut grad).and(activated).for_each(softsign_grad);
                }
                Ok(grad)
            }
            Activation::HardSigmoid => {
                // The slope is 1/6 on the linear segment and 0 on both saturated ends. An end
                // is exactly where the forward clamp wrote 0 or 1
                let mut grad = grad_output.clone();
                let hard_sigmoid_grad = |g: &mut f32, &a: &f32| {
                    *g *= if a > 0.0 && a < 1.0 {
                        HARD_SIGMOID_SLOPE
                    } else {
                        0.0
                    };
                };
                if activated.len() >= cheap_map_parallel_threshold() {
                    Zip::from(&mut grad)
                        .and(activated)
                        .par_for_each(hard_sigmoid_grad);
                } else {
                    Zip::from(&mut grad)
                        .and(activated)
                        .for_each(hard_sigmoid_grad);
                }
                Ok(grad)
            }
            Activation::Exponential => {
                let mut grad = grad_output.clone();
                let exponential_grad = |g: &mut f32, &a: &f32| {
                    *g *= a;
                };
                if activated.len() >= cheap_map_parallel_threshold() {
                    Zip::from(&mut grad)
                        .and(activated)
                        .par_for_each(exponential_grad);
                } else {
                    Zip::from(&mut grad)
                        .and(activated)
                        .for_each(exponential_grad);
                }
                Ok(grad)
            }
        }
    }

    /// Checks that a parameterized variant carries a usable parameter
    ///
    /// Every trainable layer's constructor calls this. An unusable slope or scale therefore
    /// fails where you build the model, not on the first forward pass. The parameter-free
    /// variants always pass.
    ///
    /// Both bounds are strict for the same reason. [`Activation::backward`] reads only the
    /// activated output, and it separates the 2 branches by the sign of that output. A slope
    /// or scale of 0 collapses the whole negative side onto `a = 0`, which erases the branch.
    /// A negative one inverts the sign, which reads the wrong branch
    ///
    /// # Returns
    ///
    /// - `Result<(), Error>` - Ok when the activation is usable
    ///
    /// # Notes
    ///
    /// An embedded [`Activation::Softmax`] accepts the default axis `-1` alone. This
    /// restriction is a deliberate simplification, not a limit of the math
    ///
    /// The host layers do not agree on the rank of the tensor they hand to the activation.
    /// [`Dense`](crate::neural_network::layers::dense::Dense) folds every leading axis of its
    /// output into 1 row axis, and it applies the activation to that rank-2 matrix. The
    /// convolution layers do not fold, and they apply the activation to the full output
    /// tensor. The recurrent layers apply the activation to 1 rank-2 state per timestep
    ///
    /// The axis `-1` names the same lane in all 3 families. Any other axis names a different
    /// tensor axis in each family. In `Dense` a non-final axis would normalize the folded
    /// rows, which are not the wanted lanes. A per-host axis rule is deferred, not
    /// impossible. Use the standalone [`Softmax`] layer for a different axis
    ///
    /// # Errors
    ///
    /// - `Error::InvalidParameter` - `LeakyReLU`'s `negative_slope` or `ELU`'s `alpha` is not
    ///   finite and greater than 0, or `Softmax`'s `axis` is not `-1`
    pub fn validate(&self) -> Result<(), Error> {
        let (name, value, reason) = match self {
            Activation::LeakyReLU { negative_slope } => (
                "negative_slope",
                *negative_slope,
                "must be finite and greater than 0 (use Activation::ReLU for 0)",
            ),
            Activation::ELU { alpha } => ("alpha", *alpha, "must be finite and greater than 0"),
            Activation::Softmax { axis } => {
                if *axis != DEFAULT_SOFTMAX_AXIS {
                    return Err(Error::invalid_parameter(
                        "axis",
                        "must be -1 in a trainable layer (use the Softmax layer for another axis)",
                    ));
                }
                return Ok(());
            }
            _ => return Ok(()),
        };

        if !value.is_finite() || value <= 0.0 {
            return Err(Error::invalid_parameter(name, reason));
        }
        Ok(())
    }
}

impl From<Linear> for Activation {
    #[inline]
    fn from(_: Linear) -> Self {
        Activation::Linear
    }
}
impl From<ReLU> for Activation {
    #[inline]
    fn from(_: ReLU) -> Self {
        Activation::ReLU
    }
}
impl From<Sigmoid> for Activation {
    #[inline]
    fn from(_: Sigmoid) -> Self {
        Activation::Sigmoid
    }
}
impl From<Tanh> for Activation {
    #[inline]
    fn from(_: Tanh) -> Self {
        Activation::Tanh
    }
}
impl From<Softmax> for Activation {
    #[inline]
    fn from(layer: Softmax) -> Self {
        Activation::Softmax { axis: layer.axis }
    }
}
impl From<LeakyReLU> for Activation {
    #[inline]
    fn from(layer: LeakyReLU) -> Self {
        Activation::LeakyReLU {
            negative_slope: layer.negative_slope,
        }
    }
}
impl From<ELU> for Activation {
    #[inline]
    fn from(layer: ELU) -> Self {
        Activation::ELU { alpha: layer.alpha }
    }
}
impl From<SELU> for Activation {
    #[inline]
    fn from(_: SELU) -> Self {
        Activation::SELU
    }
}
impl From<Softplus> for Activation {
    #[inline]
    fn from(_: Softplus) -> Self {
        Activation::Softplus
    }
}
impl From<Softsign> for Activation {
    #[inline]
    fn from(_: Softsign) -> Self {
        Activation::Softsign
    }
}
impl From<HardSigmoid> for Activation {
    #[inline]
    fn from(_: HardSigmoid) -> Self {
        Activation::HardSigmoid
    }
}
impl From<Exponential> for Activation {
    #[inline]
    fn from(_: Exponential) -> Self {
        Activation::Exponential
    }
}

/// Resolves a softmax axis against the rank of the tensor it applies to
///
/// A negative axis counts back from the end. The resolution happens on each call, never when
/// the [`Activation`] value or the [`Softmax`] layer is built. A value that holds a resolved
/// index would reduce the wrong axis as soon as the rank of the input changed
///
/// # Parameters
///
/// - `axis` - Axis to normalize, which can be negative
/// - `ndim` - Rank of the tensor
///
/// # Returns
///
/// - `Result<usize, Error>` - The resolved axis, from 0 through `ndim - 1`
///
/// # Errors
///
/// - `Error::InvalidInput` - The resolved axis falls outside the rank
fn resolve_softmax_axis(axis: i32, ndim: usize) -> Result<usize, Error> {
    // The widening to i64 keeps the sum in range for every i32 input, `i32::MIN` included
    let resolved = if axis < 0 {
        axis as i64 + ndim as i64
    } else {
        axis as i64
    };

    if resolved < 0 || resolved >= ndim as i64 {
        return Err(Error::invalid_input(format!(
            "Softmax axis {axis} is out of bounds for an input of rank {ndim}"
        )));
    }
    Ok(resolved as usize)
}

/// Softmax forward over 1 axis, with the lane-max shift for numerical stability
///
/// The last axis takes the fold-to-2D path, whose lanes are contiguous rows. Every other axis
/// takes the lane path, because the lanes along a non-final axis are strided and a fold would
/// mix them
fn softmax_forward(input: &Tensor, axis: i32) -> Result<Tensor, Error> {
    let shape = input.shape();
    let ndim = shape.len();
    let axis = resolve_softmax_axis(axis, ndim)?;

    let apply_softmax = |mut lane: ArrayViewMut1<f32>| {
        // Subtract the lane max so every exp argument is <= 0 (no overflow)
        let max_val = lane.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        lane.map_inplace(|x| *x = (*x - max_val).exp());
        // The max-shift guarantees 1 of the terms is exp(0)=1.0, so the sum is always >= 1.0
        let sum = lane.sum();
        lane.map_inplace(|x| *x /= sum);
    };

    if axis + 1 == ndim {
        // Flatten to [batch, features]. A fold of the leading axes keeps every last-axis lane
        // whole, so the 2D rows are exactly the lanes
        let batch_size: usize = shape[..ndim - 1].iter().product();
        let num_features = shape[ndim - 1];

        // `to_owned` keeps a non-C-order array's strides, so `into_shape_with_order` then
        // refuses it. `as_standard_layout` puts the array in C order first.
        let mut output_2d = input
            .as_standard_layout()
            .into_owned()
            .into_shape_with_order((batch_size, num_features))
            .context("Failed to reshape for softmax computation")?;

        // The task is 1 row, so a single-row input has nothing to spread and must stay serial
        if batch_size > 1 && batch_size * num_features >= exp_map_parallel_threshold() {
            output_2d
                .axis_iter_mut(Axis(0))
                .into_par_iter()
                .for_each(apply_softmax);
        } else {
            output_2d.axis_iter_mut(Axis(0)).for_each(apply_softmax);
        }

        return Ok(output_2d
            .into_shape_with_order(shape)
            .context("Failed to reshape back after softmax computation")?
            .into_dyn());
    }

    // A fresh C-order destination takes the input values, and the lanes then hold the strided
    // groups that the chosen axis selects. This keeps the result in C order for every input
    // layout, which no reshape of a strided view can promise
    let mut output = Tensor::zeros(input.raw_dim());
    output.assign(input);

    let lane_len = shape[axis];
    let lane_count: usize = shape
        .iter()
        .enumerate()
        .filter(|(index, _)| *index != axis)
        .map(|(_, &length)| length)
        .product();

    // The task is 1 lane, so a single-lane input has nothing to spread and must stay serial
    if lane_count > 1 && lane_count * lane_len >= exp_map_parallel_threshold() {
        Zip::from(output.lanes_mut(Axis(axis))).par_for_each(apply_softmax);
    } else {
        Zip::from(output.lanes_mut(Axis(axis))).for_each(apply_softmax);
    }

    Ok(output)
}

/// Softmax backward using the Jacobian-vector product expressed via the cached output
///
/// The product `dL/dz = a * (g - sum_over_axis(a * g))` is exact for every axis, so the
/// output-only derivative contract holds as long as the value carries the axis
fn softmax_backward(output: &Tensor, grad_output: &Tensor, axis: i32) -> Result<Tensor, Error> {
    let shape = output.shape();
    let ndim = shape.len();
    let axis = resolve_softmax_axis(axis, ndim)?;

    // Softmax keeps the shape, so the 2 tensors walk the same lanes. The lane path pairs the
    // 2 shapes directly, and a mismatch there is a panic rather than an error
    if grad_output.shape() != shape {
        return Err(Error::shape_mismatch(shape, grad_output.shape()));
    }

    // grad_input[i] = a[i] * (grad_output[i] - sum_j(a[j] * grad_output[j]))
    let compute_gradient = |mut grad_lane: ArrayViewMut1<f32>,
                            out_lane: ArrayView1<f32>,
                            grad_out_lane: ArrayView1<f32>| {
        let dot: f32 = out_lane
            .iter()
            .zip(grad_out_lane.iter())
            .map(|(&o, &g)| o * g)
            .sum();

        Zip::from(&mut grad_lane)
            .and(&out_lane)
            .and(&grad_out_lane)
            .for_each(|grad, &o, &g| *grad = o * (g - dot));
    };

    if axis + 1 == ndim {
        let batch_size: usize = shape[..ndim - 1].iter().product();
        let num_features = shape[ndim - 1];

        let output_2d = output
            .to_shape((batch_size, num_features))
            .context("Failed to reshape output for backward")?;

        let grad_output_2d = grad_output
            .to_shape((batch_size, num_features))
            .context("Failed to reshape grad_output for backward")?;

        let mut grad_input_2d = Array2::<f32>::zeros((batch_size, num_features));

        // The backward pass is a row dot and a scale, with no `exp`, so it is a cheap map. The
        // task is 1 row, so a single-row input stays serial
        if batch_size > 1 && batch_size * num_features >= cheap_map_parallel_threshold() {
            Zip::from(grad_input_2d.axis_iter_mut(Axis(0)))
                .and(output_2d.axis_iter(Axis(0)))
                .and(grad_output_2d.axis_iter(Axis(0)))
                .par_for_each(compute_gradient);
        } else {
            Zip::from(grad_input_2d.axis_iter_mut(Axis(0)))
                .and(output_2d.axis_iter(Axis(0)))
                .and(grad_output_2d.axis_iter(Axis(0)))
                .for_each(compute_gradient);
        }

        return Ok(grad_input_2d
            .into_shape_with_order(shape)
            .context("Failed to reshape grad_input back")?
            .into_dyn());
    }

    let mut grad_input = Tensor::zeros(output.raw_dim());

    let lane_len = shape[axis];
    let lane_count: usize = shape
        .iter()
        .enumerate()
        .filter(|(index, _)| *index != axis)
        .map(|(_, &length)| length)
        .product();

    // The task is 1 lane, so a single-lane input stays serial
    if lane_count > 1 && lane_count * lane_len >= cheap_map_parallel_threshold() {
        Zip::from(grad_input.lanes_mut(Axis(axis)))
            .and(output.lanes(Axis(axis)))
            .and(grad_output.lanes(Axis(axis)))
            .par_for_each(compute_gradient);
    } else {
        Zip::from(grad_input.lanes_mut(Axis(axis)))
            .and(output.lanes(Axis(axis)))
            .and(grad_output.lanes(Axis(axis)))
            .for_each(compute_gradient);
    }

    Ok(grad_input)
}

/// Unit tests for the activation layer helpers and the `Activation` enum
#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_abs_diff_eq;
    use ndarray::Array2;

    // Helpers

    /// Build a 2-D Tensor (ArrayD<f32>) from a row-major `data` vec with shape `(rows, cols)`
    fn tensor2(rows: usize, cols: usize, data: Vec<f32>) -> Tensor {
        Array2::from_shape_vec((rows, cols), data)
            .expect("shape/data mismatch")
            .into_dyn()
    }

    // softmax_forward

    /// Softmax of a single row matches the hand-computed distribution
    #[test]
    fn softmax_forward_basic_row() {
        let input = tensor2(1, 3, vec![0.0, 1.0, 2.0]);
        let output = softmax_forward(&input, DEFAULT_SOFTMAX_AXIS).expect("softmax_forward failed");
        let vals = output.as_slice().expect("not contiguous");
        assert_abs_diff_eq!(vals[0], 0.09003_f32, epsilon = 1e-4);
        assert_abs_diff_eq!(vals[1], 0.24473_f32, epsilon = 1e-4);
        assert_abs_diff_eq!(vals[2], 0.66524_f32, epsilon = 1e-4);
    }

    /// Softmax outputs sum to 1.0
    #[test]
    fn softmax_forward_sums_to_one() {
        let input = tensor2(1, 3, vec![0.0, 1.0, 2.0]);
        let output = softmax_forward(&input, DEFAULT_SOFTMAX_AXIS).expect("softmax_forward failed");
        let sum: f32 = output.iter().sum();
        assert_abs_diff_eq!(sum, 1.0_f32, epsilon = 1e-6);
    }

    /// A row of equal large values stays numerically stable and produces a uniform distribution
    #[test]
    fn softmax_forward_large_equal_values_stable() {
        let input = tensor2(1, 3, vec![1000.0, 1000.0, 1000.0]);
        let output = softmax_forward(&input, DEFAULT_SOFTMAX_AXIS).expect("softmax_forward failed");
        let vals = output.as_slice().expect("not contiguous");
        let third = 1.0_f32 / 3.0;
        assert_abs_diff_eq!(vals[0], third, epsilon = 1e-6);
        assert_abs_diff_eq!(vals[1], third, epsilon = 1e-6);
        assert_abs_diff_eq!(vals[2], third, epsilon = 1e-6);
    }

    /// A single-element row maps to 1.0
    #[test]
    fn softmax_forward_single_element_row() {
        let input = tensor2(1, 1, vec![5.0]);
        let output = softmax_forward(&input, DEFAULT_SOFTMAX_AXIS).expect("softmax_forward failed");
        let vals = output.as_slice().expect("not contiguous");
        assert_abs_diff_eq!(vals[0], 1.0_f32, epsilon = 1e-6);
    }

    /// An input that is not in C order gives the same result as the same values in C order
    ///
    /// `Permute` produces C order on purpose, but a caller can hand a transposed tensor
    /// straight to this layer. Every other layer accepts one
    #[test]
    fn softmax_forward_accepts_input_that_is_not_in_c_order() {
        use ndarray::IxDyn;

        // `permuted_axes` reorders the strides only, and `to_owned` keeps them. The result
        // owns a contiguous buffer while `is_standard_layout` stays false
        let base = tensor2(2, 3, vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0]);
        let transposed = base.view().permuted_axes(IxDyn(&[1, 0])).to_owned();
        assert!(
            !transposed.is_standard_layout(),
            "the test input must not be in C order"
        );

        let output = softmax_forward(&transposed, DEFAULT_SOFTMAX_AXIS)
            .expect("softmax_forward must accept it");
        assert_eq!(output.shape(), &[3, 2]);

        let c_order: Tensor = transposed.as_standard_layout().into_owned();
        let want = softmax_forward(&c_order, DEFAULT_SOFTMAX_AXIS).expect("softmax_forward failed");
        for (got, expected) in output.iter().zip(want.iter()) {
            assert_abs_diff_eq!(*got, *expected, epsilon = 1e-6);
        }
    }

    /// A 1-D input normalizes its single axis
    #[test]
    fn softmax_forward_accepts_1d_input() {
        use ndarray::Array1;
        let input = Array1::from_vec(vec![0.0_f32, 1.0, 2.0]).into_dyn();
        let output =
            softmax_forward(&input, DEFAULT_SOFTMAX_AXIS).expect("1-D input must be accepted");
        assert_eq!(output.shape(), &[3]);
        let vals: Vec<f32> = output.iter().cloned().collect();
        assert_abs_diff_eq!(vals[0], 0.09003_f32, epsilon = 1e-4);
        assert_abs_diff_eq!(vals[1], 0.24473_f32, epsilon = 1e-4);
        assert_abs_diff_eq!(vals[2], 0.66524_f32, epsilon = 1e-4);
    }

    /// A 0-D input has no axis to normalize, so every axis is out of bounds
    #[test]
    fn softmax_forward_rejects_0d_input() {
        let input: Tensor = ndarray::arr0(1.0_f32).into_dyn();
        assert!(
            matches!(
                softmax_forward(&input, DEFAULT_SOFTMAX_AXIS),
                Err(Error::InvalidInput(_))
            ),
            "0-D input must return InvalidInput"
        );
    }

    /// An axis outside the rank of the input fails the call, at both ends
    #[test]
    fn softmax_forward_rejects_out_of_range_axis() {
        let input = tensor2(2, 3, vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0]);
        for axis in [2, 7, -3, -5, i32::MIN, i32::MAX] {
            assert!(
                matches!(softmax_forward(&input, axis), Err(Error::InvalidInput(_))),
                "axis {axis} must be rejected for a rank-2 input"
            );
        }
    }

    /// The axis resolves against the rank of the input, on every call
    ///
    /// 1 value normalizes the last axis of a rank-2 input and the middle axis of a rank-3
    /// input. A resolved index held from an earlier call would reduce the wrong axis
    #[test]
    fn softmax_forward_resolves_a_negative_axis_at_call_time() {
        use ndarray::{Array, Ix3};

        let rank_2 = tensor2(2, 3, vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0]);
        let last = softmax_forward(&rank_2, -1).expect("rank-2 forward");
        let same = softmax_forward(&rank_2, 1).expect("rank-2 forward");
        assert_eq!(
            last.iter().cloned().collect::<Vec<f32>>(),
            same.iter().cloned().collect::<Vec<f32>>(),
            "-1 and 1 name the same axis of a rank-2 input"
        );

        let rank_3: Tensor =
            Array::from_shape_fn((2, 3, 4), |(i, j, k)| (i * 12 + j * 4 + k) as f32 * 0.25)
                .into_dyn();
        let middle = softmax_forward(&rank_3, -2).expect("rank-3 forward");
        let by_index = softmax_forward(&rank_3, 1).expect("rank-3 forward");
        assert_eq!(
            middle.iter().cloned().collect::<Vec<f32>>(),
            by_index.iter().cloned().collect::<Vec<f32>>(),
            "-2 and 1 name the same axis of a rank-3 input"
        );

        // Every lane of the resolved axis sums to 1, and no other axis does
        let view = middle
            .view()
            .into_dimensionality::<Ix3>()
            .expect("the result is rank 3");
        for lane in view.lanes(Axis(1)) {
            assert_abs_diff_eq!(lane.sum(), 1.0_f32, epsilon = 1e-6);
        }
    }

    /// The lane path and the fold-to-2D path give the same values for the same lanes
    ///
    /// A softmax over axis 0 of a rank-2 input runs the lane path over strided lanes. The same
    /// values transposed, normalized over the last axis, run the fold path. The 2 results must
    /// agree, which is what makes the strided lane walk trustworthy
    #[test]
    fn softmax_forward_lane_path_agrees_with_the_fold_path() {
        use ndarray::IxDyn;

        let base = tensor2(3, 4, (0..12).map(|v| v as f32 * 0.3 - 1.5).collect());
        let by_lane = softmax_forward(&base, 0).expect("axis 0 forward");

        let transposed: Tensor = {
            let view = base.view().permuted_axes(IxDyn(&[1, 0]));
            let mut owned = Tensor::zeros(view.raw_dim());
            owned.assign(&view);
            owned
        };
        let by_fold = softmax_forward(&transposed, -1).expect("axis -1 forward");

        for (row, column) in (0..3).flat_map(|r| (0..4).map(move |c| (r, c))) {
            assert_abs_diff_eq!(
                by_lane[[row, column]],
                by_fold[[column, row]],
                epsilon = 1e-7
            );
        }
    }

    // softmax_backward

    /// Backward gradient matches the hand-computed Jacobian-vector product
    #[test]
    fn softmax_backward_jacobian_vector_product() {
        let output = tensor2(1, 3, vec![0.25, 0.25, 0.5]);
        let grad_output = tensor2(1, 3, vec![1.0, 0.0, 0.0]);
        let grad_input = softmax_backward(&output, &grad_output, DEFAULT_SOFTMAX_AXIS)
            .expect("softmax_backward failed");
        let vals = grad_input.as_slice().expect("not contiguous");
        assert_abs_diff_eq!(vals[0], 0.1875_f32, epsilon = 1e-6);
        assert_abs_diff_eq!(vals[1], -0.0625_f32, epsilon = 1e-6);
        assert_abs_diff_eq!(vals[2], -0.125_f32, epsilon = 1e-6);
    }

    /// The gradient row sums to about zero, since the softmax Jacobian rows sum to zero
    #[test]
    fn softmax_backward_row_sums_to_zero() {
        let output = tensor2(1, 3, vec![0.25, 0.25, 0.5]);
        let grad_output = tensor2(1, 3, vec![1.0, 0.0, 0.0]);
        let grad_input = softmax_backward(&output, &grad_output, DEFAULT_SOFTMAX_AXIS)
            .expect("softmax_backward failed");
        let row_sum: f32 = grad_input.iter().sum();
        assert_abs_diff_eq!(row_sum, 0.0_f32, epsilon = 1e-6);
    }

    /// A grad_output that does not match the output shape is an error, not a panic
    ///
    /// The lane path pairs the 2 shapes with a `Zip`, which panics on a mismatch. The guard
    /// must run before that, so the caller gets `Error::ShapeMismatch`. The fold path takes
    /// the same guard, and this test pins both paths
    #[test]
    fn softmax_backward_rejects_a_mismatched_grad_output() {
        use ndarray::IxDyn;

        let output = Tensor::zeros(IxDyn(&[2, 3, 4]));
        let grad_output = Tensor::zeros(IxDyn(&[2, 3, 5]));

        // Axis 1 is not the last axis, so this call takes the lane path
        match softmax_backward(&output, &grad_output, 1) {
            Err(Error::ShapeMismatch { expected, found }) => {
                assert_eq!(expected, vec![2, 3, 4]);
                assert_eq!(found, vec![2, 3, 5]);
            }
            other => panic!("the lane path must return ShapeMismatch, got {other:?}"),
        }

        assert!(
            matches!(
                softmax_backward(&output, &grad_output, DEFAULT_SOFTMAX_AXIS),
                Err(Error::ShapeMismatch { .. })
            ),
            "the fold path must return ShapeMismatch as well"
        );
    }

    // Round-trip tests for the Activation enum's public API

    /// Activation::Softmax forward delegates to softmax_forward
    #[test]
    fn activation_softmax_forward_via_enum() {
        let input = tensor2(1, 3, vec![0.0, 1.0, 2.0]);
        let output = Activation::Softmax {
            axis: DEFAULT_SOFTMAX_AXIS,
        }
        .forward(&input)
        .expect("Activation::Softmax forward failed");
        let vals = output.as_slice().expect("not contiguous");
        // Same expected values as softmax_forward_basic_row
        assert_abs_diff_eq!(vals[0], 0.09003_f32, epsilon = 1e-4);
        assert_abs_diff_eq!(vals[1], 0.24473_f32, epsilon = 1e-4);
        assert_abs_diff_eq!(vals[2], 0.66524_f32, epsilon = 1e-4);
    }

    /// Activation::Softmax backward delegates to softmax_backward
    #[test]
    fn activation_softmax_backward_via_enum() {
        let output = tensor2(1, 3, vec![0.25, 0.25, 0.5]);
        let grad_output = tensor2(1, 3, vec![1.0, 0.0, 0.0]);
        let grad_input = Activation::Softmax {
            axis: DEFAULT_SOFTMAX_AXIS,
        }
        .backward(&output, &grad_output)
        .expect("Activation::Softmax backward failed");
        let vals = grad_input.as_slice().expect("not contiguous");
        assert_abs_diff_eq!(vals[0], 0.1875_f32, epsilon = 1e-6);
        assert_abs_diff_eq!(vals[1], -0.0625_f32, epsilon = 1e-6);
        assert_abs_diff_eq!(vals[2], -0.125_f32, epsilon = 1e-6);
    }

    // The pinned float32 tables below confirm agreement with reference values, not just
    // that `forward` and `backward` match each other

    /// The probe inputs shared by every pinned table below
    const PROBES: [f32; 11] = [-5.0, -3.0, -2.0, -1.0, -0.5, 0.0, 0.5, 1.0, 2.0, 3.0, 5.0];

    /// Asserts that `actual` matches the pinned `expected` row
    ///
    /// The tolerance mixes a relative and an absolute bound, because the tables span 5
    /// orders of magnitude. `Exponential` reaches 148, where 1 f32 ulp is already about 1e-5
    fn assert_pinned(name: &str, actual: &Tensor, expected: &[f32]) {
        let got: Vec<f32> = actual.iter().cloned().collect();
        assert_eq!(got.len(), expected.len(), "{name}: length mismatch");
        for (i, (&g, &e)) in got.iter().zip(expected.iter()).enumerate() {
            let tol = 1e-5 * e.abs().max(1.0);
            assert!(
                (g - e).abs() <= tol,
                "{name}[{i}] at x = {}: got {g}, want {e}, tolerance {tol}",
                PROBES[i]
            );
        }
    }

    /// Runs `activation` over [`PROBES`] and checks both the output and the derivative
    ///
    /// An all-ones upstream gradient makes the backward result the derivative itself
    fn check_against_reference(name: &str, activation: Activation, fwd: &[f32], grad: &[f32]) {
        let input = tensor2(1, PROBES.len(), PROBES.to_vec());
        let output = activation.forward(&input).expect("forward failed");
        assert_pinned(&format!("{name} forward"), &output, fwd);

        let ones = tensor2(1, PROBES.len(), vec![1.0; PROBES.len()]);
        let derivative = activation
            .backward(&output, &ones)
            .expect("backward failed");
        assert_pinned(&format!("{name} backward"), &derivative, grad);
    }

    /// LeakyReLU with the layer's default slope of 0.3. The derivative at exactly 0 is 1,
    /// because the positive branch is `x >= 0`
    #[test]
    fn leaky_relu_matches_reference() {
        check_against_reference(
            "LeakyReLU(0.3)",
            Activation::LeakyReLU {
                negative_slope: 0.3,
            },
            &[
                -1.5,
                -0.90000004,
                -0.6,
                -0.3,
                -0.15,
                0.0,
                0.5,
                1.0,
                2.0,
                3.0,
                5.0,
            ],
            &[0.3, 0.3, 0.3, 0.3, 0.3, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
        );
    }

    /// ELU with the default alpha of 1.0
    #[test]
    fn elu_matches_reference() {
        check_against_reference(
            "ELU(1.0)",
            Activation::ELU { alpha: 1.0 },
            &[
                -0.99326205,
                -0.95021296,
                -0.86466473,
                -0.63212055,
                -0.39346933,
                0.0,
                0.5,
                1.0,
                2.0,
                3.0,
                5.0,
            ],
            &[
                0.0067379475,
                0.049787045,
                0.13533527,
                0.36787945,
                0.60653067,
                1.0,
                1.0,
                1.0,
                1.0,
                1.0,
                1.0,
            ],
        );
    }

    /// ELU with alpha = 0.5, the case that pins the branch convention
    ///
    /// At exactly 0 the derivative is `alpha`, not 1, because the positive branch is `x > 0`.
    /// The default alpha of 1.0 hides that, since both branches then give 1
    #[test]
    fn elu_half_alpha_matches_reference() {
        check_against_reference(
            "ELU(0.5)",
            Activation::ELU { alpha: 0.5 },
            &[
                -0.49663103,
                -0.47510648,
                -0.43233237,
                -0.31606027,
                -0.19673467,
                0.0,
                0.5,
                1.0,
                2.0,
                3.0,
                5.0,
            ],
            &[
                0.0033689737,
                0.024893522,
                0.067_667_63,
                0.18393973,
                0.30326533,
                0.5,
                1.0,
                1.0,
                1.0,
                1.0,
                1.0,
            ],
        );
    }

    /// SELU. The derivative at exactly 0 is `scale * alpha`, the same `x > 0` convention as ELU
    #[test]
    fn selu_matches_reference() {
        check_against_reference(
            "SELU",
            Activation::SELU,
            &[
                -1.7462534,
                -1.6705688,
                -1.5201665,
                -1.1113307,
                -0.691_758_2,
                0.0,
                0.525_350_5,
                1.050701,
                2.101402,
                3.152_103,
                5.253_505,
            ],
            &[
                0.011845981,
                0.087_530_57,
                0.23793285,
                0.646_768_6,
                1.0663412,
                1.7580993,
                1.050701,
                1.050701,
                1.050701,
                1.050701,
                1.050701,
            ],
        );
    }

    /// Softplus. The value at 0 is ln(2), and the derivative there is 0.5
    #[test]
    fn softplus_matches_reference() {
        check_against_reference(
            "Softplus",
            Activation::Softplus,
            &[
                0.0067153485,
                0.048587352,
                0.126928,
                0.313_261_7,
                0.474_077,
                // softplus(0) = ln(1 + 1), so the pinned value is exactly this constant
                std::f32::consts::LN_2,
                0.974_077,
                1.3132617,
                2.126_928,
                3.0485873,
                5.0067153,
            ],
            &[
                0.006692851,
                0.047425874,
                0.11920291,
                0.2689414,
                0.37754068,
                0.5,
                0.62245935,
                0.73105854,
                0.880_797,
                0.95257413,
                0.993_307_2,
            ],
        );
    }

    /// Softsign
    #[test]
    fn softsign_matches_reference() {
        check_against_reference(
            "Softsign",
            Activation::Softsign,
            &[
                -0.833_333_3,
                -0.75,
                -0.666_666_7,
                -0.5,
                -0.33333334,
                0.0,
                0.33333334,
                0.5,
                0.666_666_7,
                0.75,
                0.833_333_3,
            ],
            &[
                0.027777776,
                0.0625,
                0.11111112,
                0.25,
                0.44444448,
                1.0,
                0.44444448,
                0.25,
                0.11111112,
                0.0625,
                0.027777776,
            ],
        );
    }

    /// HardSigmoid. Both saturated ends are exact, and their derivative is 0
    #[test]
    fn hard_sigmoid_matches_reference() {
        check_against_reference(
            "HardSigmoid",
            Activation::HardSigmoid,
            &[
                0.0,
                0.0,
                0.16666667,
                0.33333334,
                0.416_666_7,
                0.5,
                0.583_333_4,
                0.666_666_7,
                0.833_333_4,
                1.0,
                1.0,
            ],
            &[
                0.0, 0.0, 0.16666667, 0.16666667, 0.16666667, 0.16666667, 0.16666667, 0.16666667,
                0.16666667, 0.0, 0.0,
            ],
        );
    }

    /// Exponential, whose derivative equals its output
    #[test]
    fn exponential_matches_reference() {
        let table = [
            0.006737947,
            0.049787067,
            0.13533528,
            0.36787945,
            0.60653067,
            1.0,
            1.6487212,
            2.7182817,
            7.389_056,
            20.085537,
            148.41316,
        ];
        check_against_reference("Exponential", Activation::Exponential, &table, &table);
    }

    /// The softplus derivative survives the far negative tail
    ///
    /// At x = -40 the derivative is about 4.25e-18. The direct form `1 - e^-a` rounds to
    /// exactly 0, because `e^-a` is 1 at f32 precision. `exp_m1` keeps the value
    #[test]
    fn softplus_backward_keeps_the_far_negative_tail() {
        let input = tensor2(1, 1, vec![-40.0]);
        let output = Activation::Softplus.forward(&input).expect("forward");
        let ones = tensor2(1, 1, vec![1.0]);
        let derivative = Activation::Softplus
            .backward(&output, &ones)
            .expect("backward");

        let got = derivative.iter().next().copied().expect("1 element");
        let expected = 4.248_354e-18_f32;
        assert!(
            (got - expected).abs() <= 1e-5 * expected,
            "softplus derivative at x = -40: got {got}, want {expected}"
        );
    }

    /// `validate` rejects a slope of 0 or below, since the backward pass reads the branch
    /// off the sign of the output
    #[test]
    fn validate_rejects_unusable_leaky_relu_slope() {
        for slope in [0.0, -0.1, f32::NAN, f32::INFINITY] {
            let result = Activation::LeakyReLU {
                negative_slope: slope,
            }
            .validate();
            assert!(
                matches!(result, Err(Error::InvalidParameter { .. })),
                "slope {slope} must be rejected, got {result:?}"
            );
        }
    }

    /// The same bound applies to ELU's alpha
    #[test]
    fn validate_rejects_unusable_elu_alpha() {
        for alpha in [0.0, -1.0, f32::NAN, f32::NEG_INFINITY] {
            let result = Activation::ELU { alpha }.validate();
            assert!(
                matches!(result, Err(Error::InvalidParameter { .. })),
                "alpha {alpha} must be rejected, got {result:?}"
            );
        }
    }

    /// Every usable activation passes validation
    #[test]
    fn validate_accepts_usable_activations() {
        let usable = [
            Activation::Linear,
            Activation::ReLU,
            Activation::Sigmoid,
            Activation::Tanh,
            Activation::Softmax {
                axis: DEFAULT_SOFTMAX_AXIS,
            },
            Activation::LeakyReLU {
                negative_slope: 0.3,
            },
            Activation::ELU { alpha: 1.0 },
            Activation::SELU,
            Activation::Softplus,
            Activation::Softsign,
            Activation::HardSigmoid,
            Activation::Exponential,
        ];
        for activation in usable {
            assert!(
                activation.validate().is_ok(),
                "{activation:?} must pass validation"
            );
        }
    }
}