3.2. Dense Layers and Activations

Dense is the main layer of a feedforward network. It applies a linear map, input * W + b, then an elementwise nonlinearity. RustyML fuses that nonlinearity into the layer itself. It does not treat the nonlinearity as a separate stage.

This page covers the constructor and its fused-activation design. It also covers the exact initialization scheme and the 12 activations, with their gradients and failure modes. It covers the GEMM-backed compute path and how to read or inject weights. It assumes you have already read the Sequential model page. The layers here are what you collect in a SequentialBuilder, and the input width of each one comes from the model build.

3.2.1. The fused-activation design

The constructor takes the activation as its second argument, not as a separate layer:

pub fn new(units: usize, activation: impl Into<Activation>) -> Result<Dense, Error>

Internally the layer stores an Activation value. This is a plain Copy enum with 12 variants. It is not a generic type parameter Dense<A>, and it is not a Box<dyn Activation>. That choice is deliberate. A generic parameter would monomorphize Dense 12 ways.

It would also force weight deserialization to probe every Dense<A> pairing to find the concrete type on disk. A trait object would add an indirection to every elementwise call. The runtime enum keeps Dense a single concrete type. The persistence layer can then downcast every saved layer to exactly 1 struct (see Saving and Loading Weights). The activation math stays a pure, stateless function, and the layer calls it inside its own forward and backward passes.

Keras makes the same fusion choice, with Dense(units, activation="relu"). RustyML has no implicit “no activation”. You must always pass one. To get a pure linear layer, pass Activation::Linear, the identity. This is the equivalent of Keras’ activation=None. Every example below that ends in a regression head does exactly this.

Fusing is the right default. RustyML also provides every activation as a standalone layer (Linear, ReLU, LeakyReLU, ELU, SELU, Sigmoid, HardSigmoid, Tanh, Softplus, Softsign, Exponential, Softmax). Use these when you need something between the linear map and the nonlinearity. A normalization layer is the most common case (see Regularization and Normalization Layers). The 2 forms are equivalent:

// These two stacks compute the same thing.
builder.add(Dense::new(32, Activation::ReLU).unwrap());        // fused: one cached output

builder.add(Dense::new(32, Activation::Linear).unwrap())       // split: identity Dense ...
       .add(ReLU::new());                                      // ... then a standalone ReLU

Prefer the fused form. It caches a single activated tensor instead of 2. The backward pass then differentiates the activation in terms of that cached output (see 3.2.6). Use the split form only when a layer must sit between W*x + b and the nonlinearity.

3.2.2. Constructing and sizing a layer

Dense::new(units, activation) returns Result<Dense, Error>. units is the number of neurons, and therefore the output width. It must be non-zero, and a 0 yields Error::InvalidParameter (see Error Handling).

The constructor no longer takes the input width. A kernel extent comes from the input, and the input is not there when you write the constructor. Layer::build is where the layer learns it. build reads the last axis of the shape it receives, records it, and draws the kernel at that width. SequentialBuilder::build calls it once per layer, so a stack of Dense layers never repeats a width. 2 adjacent layers can no longer disagree about one.

The parameter count is input_width * units + units: 1 weight per (input, output) pair, plus 1 bias per output. An unbuilt layer holds no array, so it reports 0 for both counts and output_shape() reports "Unknown". param_count() reports it as ParamCounts::trainable(n), which is n trainable elements and 0 non-trainable ones.

The layer gives its 2 tensors the Keras names kernel and bias. The optimizer keys its state on those names, and a checkpoint addresses them by the same names (see 3.4 and 3.9). output_shape() renders (None, units) once the layer is built for a free batch axis, mirroring Keras’ summary. A Dense(3, ...) built for a last axis of 4 therefore holds a 4 x 3 weight matrix and a 1 x 3 bias. That is 12 + 3 = 15 trainable scalars.

compute_output_shape answers the same question without any of that state. It is pure: it replaces the last axis of the shape you hand it with units, and it reads no cache. It works on a layer that was never built and never run. A layer that is built also checks the last axis against the width it holds, and refuses a shape that disagrees.

with_use_bias(false) builds the layer without the bias, exactly as the Keras use_bias argument does. The forward pass is then activation(input * W), the parameter count drops to input_width * units, and the layer holds 1 array in place of 2. Put the flag on a layer that a normalization layer follows. The shift of that layer makes the bias of the dense layer redundant.

Input is a tensor of rank 2 or more, whose last axis holds the built width. The layer contracts that last axis alone and replaces it with units. Every axis in front of it passes through untouched, so a (batch, features) input gives (batch, units), and a (batch, timesteps, features) input gives (batch, timesteps, units). A rank-1 input, and a last axis of the wrong width, are both rejected with Error::InvalidInput rather than silently reshaped. The following program constructs a layer, builds it, inspects it, and injects known weights:

use ndarray::Array2;
use rustyml::neural_network::Shape;
use rustyml::neural_network::layers::{Activation, Dense, ParamCounts};
use rustyml::neural_network::traits::{Layer, LayerBase, UnaryLayer};

fn main() {
    // The constructor takes the unit count alone, and it draws nothing.
    let mut dense = Dense::new(3, Activation::ReLU).unwrap();
    assert_eq!(dense.param_count(), ParamCounts::trainable(0));
    println!("before build: {}", dense.output_shape()); // Unknown

    // `build` supplies the input shape. The layer reads 4 off the last axis and allocates.
    dense.build(&Shape::with_free_batch(&[8, 4])).unwrap();
    assert_eq!(dense.param_count(), ParamCounts::trainable(15)); // 4*3 + 3
    println!("after build:  {}", dense.output_shape()); // (None, 3)

    // Read the drawn arrays by name, with no clone.
    for entry in dense.weights() {
        println!("{} {:?} {:?}", entry.name, entry.kind, entry.value.shape());
    }
    assert_eq!(dense.weight("kernel").unwrap().shape(), [4, 3]);
    assert_eq!(dense.weight("bias").unwrap().shape(), [1, 3]);

    // Inject known weights. The shapes are validated against what the build settled.
    let weights =
        Array2::from_shape_vec((4, 3), (0..12).map(|v| v as f32).collect::<Vec<f32>>()).unwrap();
    let bias = Array2::zeros((1, 3));
    dense.set_weights(weights, bias).unwrap();

    // `with_use_bias(false)` drops the bias array, so the layer holds the kernel alone.
    let mut no_bias = Dense::new(3, Activation::ReLU).unwrap().with_use_bias(false);
    no_bias.build(&Shape::with_free_batch(&[8, 4])).unwrap();
    assert_eq!(no_bias.param_count(), ParamCounts::trainable(12));
    assert!(no_bias.weight("bias").is_none());
    println!("no-bias paths: {}", no_bias.weights().len());

    // The shape answer is pure, so an unbuilt twin answers for any rank.
    let twin = Dense::new(3, Activation::ReLU).unwrap();
    let sequence = Shape::new(vec![None, Some(5), Some(4)]);
    println!("rank 3: {}", twin.compute_output_shape(&sequence).unwrap()); // (None, 5, 3)
}

set_weights needs a built layer, and it returns Error::NeuralNetwork(NnError::NotBuilt) on a layer that holds no array yet. There is nothing to check a shape against before the build. It then checks both shapes and returns Error::NeuralNetwork(NnError::WeightShape) on a mismatch. set_weights rejects a (3, 3) weight for a 4 -> 3 layer, or a (1, 4) bias, rather than truncating it.

Its second argument is optional. Pass None for a layer built with with_use_bias(false), and pass the array for every other layer. The other order returns Error::InvalidParameter, because a layer that holds no bias would silently drop one that reaches it.

The rank costs nothing. 1 kernel serves every leading position. The layer folds all the leading axes into 1 row axis, and it runs the same single matrix product a rank-2 input runs. A rank-3 pass and the folded rank-2 pass therefore give the same values, bit for bit.

forward_mut is the entry point of a caller that drives 1 layer by hand. It builds the layer from the tensor it receives, and then it completes the pass. This is why the 2 layers below need no build call. The layer parks the shape of the input it received, so the gradient goes back at the rank it arrived at. The build shape holds the last axis alone, so output_shape() reports (None, units) at every rank. This is what lets a Dense layer sit after a recurrent layer that returns sequences, and transform every timestep with the same weights:

use ndarray::Array;
use rustyml::neural_network::Ctx;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::{ParamId, UnaryLayer};

fn main() {
    // 2 samples, 5 timesteps, 4 features per step.
    let x = Array::from_shape_fn((2, 5, 4), |(n, t, f)| (n * 20 + t * 4 + f) as f32 * 0.01)
        .into_dyn();

    // The kernel contracts the last axis, and every leading axis passes through. A
    // training context holds what the backward pass needs.
    let mut dense = Dense::new(3, Activation::Linear).unwrap().with_random_state(7);
    let mut ctx = Ctx::training();
    let out = dense.forward_mut(&x, &mut ctx).unwrap();
    assert_eq!(out.shape(), &[2, 5, 3]);

    // 1 kernel serves every leading position, so a rank-3 pass equals a folded rank-2
    // pass. An inference context parks nothing at all, so it stays empty.
    let folded = x.to_shape((10, 4)).unwrap().to_owned().into_dyn();
    let mut twin = Dense::new(3, Activation::Linear).unwrap().with_random_state(7);
    let mut eval = Ctx::inference();
    let flat = twin.forward_mut(&folded, &mut eval).unwrap();
    assert_eq!(flat.shape(), &[10, 3]);
    assert_eq!(eval.pending_caches(), 0);
    for row in 0..10 {
        for u in 0..3 {
            assert_eq!(out[[row / 5, row % 5, u]], flat[[row, u]]);
        }
    }

    // The gradient goes back at the rank of the input, and the 2 parameter gradients
    // go into the store of the context.
    let grad = dense.backward(&out, &mut ctx).unwrap();
    assert_eq!(grad.shape(), &[2, 5, 4]);
    let kernel_grad = ctx.grads().get(ParamId::new(0, "kernel")).unwrap();
    assert_eq!(kernel_grad.shape(), &[4, 3]);

    println!("output shape: {:?}", out.shape());
}

Each layer takes its own context above. A context serves 1 pass of 1 stack, and the caches of 1 position form a stack that the matching backward pass empties. 2 layers that write into 1 context under 1 position would therefore cross their caches. A model gives every layer its own position, and a caller that drives 1 layer by hand takes position 0.

A convolutional stack still needs a Flatten in front of its classifier head. Dense would otherwise contract the channel axis alone and leave every spatial position separate, which is a 1x1 convolution and not a classifier. Insert Flatten when you want 1 vector per sample.

3.2.3. Weight initialization

Weights use Xavier/Glorot uniform initialization. Each element is drawn from Uniform(-limit, +limit), where limit = sqrt(6 / (fan_in + fan_out)). For a Dense layer the 2 fans are the built input width and units. Biases start at exactly zero. It matches Keras’ Dense default exactly (glorot_uniform kernel, zeros bias).

The rule is a value, not a hidden constant. Initializer is a small closed enum with 3 variants. GlorotUniform covers almost every kernel of the module, Uniform { limit } covers the embedding table, and Orthogonal covers the recurrent kernels. A layer holds the value it draws with, so the rule compares, copies, and serializes with the layer.

The layer supplies the 2 fans, and nothing derives them from the shape of the array. They arrive as a Fans pair that the layer builds from its own configuration. The transposed convolution is the reason. A plain convolution holds its kernel as (spatial axes, channels, filters), and a transposed convolution holds the same kernel as (spatial axes, filters, channels).

The last 2 axes carry opposite roles in the 2 layouts. A rule that reads them gives the transposed layer the 2 fans the wrong way round. Glorot hides that swap, because its range reads the sum of the pair and a sum is symmetric. A layer that names its 2 counts cannot make the mistake at all. The defect never waits for an initializer that reads 1 fan alone.

RustyML applies Glorot regardless of the activation. It does not switch to He/Kaiming initialization for ReLU layers, even though He is the textbook match for rectifiers. For shallow nets this rarely matters. For a deep ReLU stack, early convergence may be slightly slower than a He-initialized equivalent. If that happens, initialize with set_weights from your own scaled draw.

Initialization draws through the crate’s shared RNG, in build. By default the seed comes from the process-global seed if one is set, otherwise from entropy. 2 runs then give different weights unless you fix the randomness. There are 2 ways to make it reproducible. Set a thread-local global seed with rustyml::random::set_global_seed(...) before you build the model. This also fixes dropout masks and the fit-time batch shuffle.

You can instead seed 1 layer explicitly, with Dense::new(...)?.with_random_state(seed). On an unbuilt layer the call only records the seed, because there is no array to redraw yet, and build spends it. On a layer that is already built the call redraws at once, so the order of the 2 calls never matters. All 15 builders that draw a weight work this way.

Either way the seed leaves the global stream untouched. The full seeding model is in Reproducibility and Random Seeds. It also explains why an explicit per-layer seed does not change the seeds handed to unseeded layers.

3.2.4. The activations

Activation has 12 variants. Each variant’s backward pass is expressed in terms of the activated output a = f(x), not the pre-activation x. The layer caches that output, not the input. That contract admits every activation whose derivative has a closed form in a, which covers all 12 below.

The same contract is exactly why GELU, SiLU (Swish), and Mish are absent. Their a = x * g(x) shape has no closed-form inverse, so no derivative in a alone exists. Adding them needs a wider contract that also hands the backward pass the pre-activation.

ReLU is the default hidden-layer choice. It is cheap, and it does not saturate on the positive side. Its failure mode is the dead neuron. The derivative is 0 for x <= 0, so a neuron whose pre-activation is negative for every example in the batch gets zero gradient.

It never updates, and it stays off for good. A high learning rate makes this worse. It pushes neurons into the dead region early. Glorot init has no leaky variant. Your options are a smaller learning rate and well-scaled inputs, or, if you inject your own weights, a better initial scale.

LeakyReLU, ELU, and SELU are the direct answer to that failure mode. Each one keeps a non-zero gradient below 0. A unit whose pre-activation stays negative for the whole batch still receives gradient, so it can recover. LeakyReLU { negative_slope } scales the negative side by a constant, and its output stays unbounded. PReLU in 3.2.5 goes further and learns that constant. ELU { alpha } instead saturates the negative side at -alpha, which pulls the mean activation toward 0, at a cost of 1 exponential per negative element.

SELU is the same shape as ELU, with alpha and scale fixed at 1.6732632 and 1.0507010. Its self-normalizing property holds only under Lecun-normal initialization, and Dense hard-codes Glorot uniform (see 3.2.3). RustyML has no Lecun-normal initializer yet. SELU still trains as a plain activation under Glorot. The variance-preserving guarantee of Klambauer et al. (2017) does not hold there. To get it, inject your own Lecun-normal draw with set_weights.

negative_slope and alpha must both be finite and greater than 0. Activation::validate enforces that bound. All 14 trainable layer constructors call it: Dense, Conv1D, Conv2D, Conv3D, Conv1DTranspose, Conv2DTranspose, Conv3DTranspose, DepthwiseConv1D, DepthwiseConv2D, SeparableConv1D, SeparableConv2D, SimpleRNN, LSTM, and GRU. An unusable value therefore returns Error::InvalidParameter where you build the model, not on the first forward pass.

The bound is strict, and the output-only contract is the reason. The backward pass reads the branch off the sign of a. A value of 0 collapses the whole negative side onto a = 0, which erases the branch. A negative value inverts the sign, so the backward pass reads the wrong branch. Use Activation::ReLU when you want a slope of 0.

LeakyReLU uses x >= 0 for the positive branch, so its derivative at exactly 0 is 1. ELU and SELU use x > 0, which puts 0 itself on the negative branch. Their derivatives at exactly 0 are therefore alpha and scale * alpha, not 1. LeakyReLU::default() uses a slope of 0.3, and ELU::default() uses an alpha of 1.0.

The following hidden stack uses ELU twice. The first layer fuses it into Dense. The second pairs an identity Dense with the standalone ELU layer, which is the split form of 3.2.1:

use ndarray::Array;
use rustyml::neural_network::Shape;
use rustyml::neural_network::layers::{Activation, Dense, ELU};
use rustyml::neural_network::losses::MeanSquaredError;
use rustyml::neural_network::optimizers::Adam;
use rustyml::neural_network::sequential::SequentialBuilder;

fn main() {
    // 4 samples, 3 features, 1 continuous target each. Negative features exercise the
    // saturating branch of ELU.
    let x = Array::from_shape_vec(
        (4, 3),
        vec![-1.0, 0.1, 0.2, 1.0, -0.9, 0.8, 0.2, 0.1, -2.0, 0.9, 1.0, 0.8],
    )
    .unwrap()
    .into_dyn();
    let y = Array::from_shape_vec((4, 1), vec![0.3, 2.7, 0.3, 2.7])
        .unwrap()
        .into_dyn();

    let mut model = SequentialBuilder::new()
        // Fused: the activation runs inside the Dense layer.
        .add(Dense::new(8, Activation::ELU { alpha: 1.0 }).unwrap())
        // Split: an identity Dense, then the same activation as its own layer.
        .add(Dense::new(8, Activation::Linear).unwrap())
        .add(ELU::new(1.0).unwrap())
        .add(Dense::new(1, Activation::Linear).unwrap()) // regression head
        .build(&Shape::known(x.shape()))
        .unwrap();
    model.compile(
        Adam::new(0.01, 0.9, 0.999, 1e-8, 0.0).unwrap(),
        MeanSquaredError::new(),
    );

    model.fit(&x, &y, 20).unwrap();

    let preds = model.predict(&x).unwrap();
    assert_eq!(preds.shape(), &[4, 1]);
    assert!(preds.iter().all(|v| v.is_finite()));

    // A slope of 0 erases the negative branch, so the constructor rejects it.
    let dead_slope = Activation::LeakyReLU { negative_slope: 0.0 };
    assert!(Dense::new(8, dead_slope).is_err());
}

Sigmoid squashes values to (0, 1). Its derivative a * (1 - a) peaks at 0.25 when a = 0.5, and it decays toward zero as the output saturates. Stacking sigmoids in a deep hidden path throttles gradients: the classic vanishing-gradient problem. Use Sigmoid for a single binary output, paired with binary cross-entropy, or as a gate. Do not use it as a deep hidden nonlinearity. At extreme inputs, f32 saturates it to exactly 0.0 or 1.0.

Tanh maps to (-1, 1). Unlike sigmoid, it is zero-centered, which tends to make hidden-layer optimization better behaved. Its gradient 1 - a^2 still reaches 1 near the origin. It saturates the same way sigmoid does at the tails. It is the natural bounded, zero-centered hidden activation, and the recurrent layers use it internally.

Softplus and Softsign are the smooth option and the bounded option. Softplus is ln(1 + e^x), a smooth approximation of ReLU. Its output is strictly positive, and its derivative 1 - e^-a never reaches 0, so it has no dead unit at all. It costs 1 exponential and 1 logarithm per element, which makes it more expensive than LeakyReLU.

Softsign is x / (1 + abs(x)). It is bounded to (-1, 1) and zero-centered, exactly like tanh. It saturates polynomially rather than exponentially, so its tails keep more gradient than the tails of Tanh. It also needs no exponential. Treat it as the cheaper, slower-saturating substitute for Tanh.

HardSigmoid approximates sigmoid with a clipped straight line, and it uses no exponential. It is clip(x/6 + 0.5, 0, 1), so it reaches exactly 0 at x = -3 and exactly 1 at x = 3. True sigmoid only approaches both ends and never arrives. The derivative is the constant 1/6 on the linear segment, and exactly 0 outside it. That flat outer gradient is the same dead-unit risk ReLU carries. Use HardSigmoid as a gate where the cost of exp matters.

Exponential is e^x. Its output is strictly positive, and its derivative is the output itself. Use it on a head that must emit a positive quantity, such as a rate, a variance, or a scale. It has no upper bound, so a pre-activation above about 88 overflows f32 to inf. Keep the head’s input scaled, or emit a log-scale value through Linear instead.

Softmax turns a lane of logits into a probability distribution. Softmax { axis } names the axis that the lanes run along, and the default is -1, the last axis. The forward pass subtracts each lane’s maximum before it takes the exponent. This makes the largest term exp(0) = 1, and the sum always >= 1. That shift makes softmax overflow-proof and shift-invariant: adding a constant to every logit leaves the output unchanged.

Softmax accepts an input of any rank. A rank-1 input is legal, and it normalizes its single axis, which matches Keras. The backward pass is the true Jacobian-vector product across the lane, not an elementwise multiply, and each gradient lane sums to zero. An axis that resolves outside the rank of the input returns Error::InvalidInput.

A negative axis counts back from the end. The value resolves against the rank of the input on each call, never when you build it. The same axis therefore means a different axis for inputs of different rank. An axis of -1 is axis 1 of a rank-2 input, and axis 3 of a rank-4 input. The standalone Softmax layer takes its axis from the with_axis builder, and Softmax::new() starts at the default -1.

An embedded softmax, the kind you pass to Dense or to a convolution layer, accepts the default axis alone. Building the layer refuses any other axis, with Error::InvalidParameter, not on the first forward pass. Dense folds every leading axis into 1 row axis, and that fold keeps each last-axis lane whole while it destroys every other axis. A convolution layer instead applies the activation to its full output tensor, with no fold. Every host family still restricts an embedded softmax to axis -1, because that axis alone names the same lane in each family. Use the standalone Softmax layer for any other axis:

use ndarray::Array;
use rustyml::neural_network::Ctx;
use rustyml::neural_network::layers::{Activation, Dense, Softmax};
use rustyml::neural_network::traits::UnaryLayer;

fn main() {
    // 2 samples, 4 positions, 3 channels.
    let x = Array::from_shape_fn((2, 4, 3), |(n, t, c)| (n * 12 + t * 3 + c) as f32 * 0.1)
        .into_dyn();

    // An inference context parks nothing, because no backward pass follows here.
    let mut ctx = Ctx::inference();

    // The default axis is -1, the last axis, so each of the 8 positions becomes a
    // distribution over the 3 channels.
    let mut over_channels = Softmax::new();
    let per_position = over_channels.forward_mut(&x, &mut ctx).unwrap();
    let first_position: f32 = (0..3).map(|c| per_position[[0, 0, c]]).sum();
    assert!((first_position - 1.0).abs() < 1e-6);

    // with_axis picks another axis. Axis 1 makes each of the 6 channel columns a
    // distribution over the 4 positions.
    let mut over_positions = Softmax::new().with_axis(1);
    let per_channel = over_positions.forward_mut(&x, &mut ctx).unwrap();
    let first_channel: f32 = (0..4).map(|t| per_channel[[0, t, 0]]).sum();
    assert!((first_channel - 1.0).abs() < 1e-6);

    // A rank-1 input normalizes its single axis.
    let logits = Array::from_vec(vec![1.0_f32, 2.0, 3.0]).into_dyn();
    let probabilities = Softmax::new().forward_mut(&logits, &mut ctx).unwrap();
    assert_eq!(probabilities.shape(), &[3]);

    // An embedded softmax accepts the default axis alone.
    assert!(Dense::new(2, Activation::Softmax { axis: -1 }).is_ok());
    assert!(Dense::new(2, Activation::Softmax { axis: 1 }).is_err());

    println!("lane sums: {first_position} and {first_channel}");
}

Softmax belongs on the output layer of a classifier, paired with the correct loss configuration. Mid-network, softmax is mechanically valid: the Jacobian backward is correct, so a Dense(..., Softmax) buried in the stack still trains. It is almost never what you want, though. It collapses the representation onto a simplex, discards all magnitude information, and saturates worse than ReLU or tanh.

Linear is the identity. Its gradient is 1, and its output is unbounded. Use it for regression heads, and whenever the loss function expects raw logits.

The loss function must match the activation on the output head. There are 2 correct pairings for multi-class classification. Mixing them silently corrupts training:

• Softmax head + CategoricalCrossEntropy::new(false). The final Dense emits probabilities. The loss consumes probabilities. Correct.
• Linear head + CategoricalCrossEntropy::new(true). The final Dense emits raw logits. The loss applies a numerically stable log-softmax internally, and returns the fused (softmax(z) - y) gradient in one step. Also correct, and more numerically stable.

The 2 pairings produce the same gradient mathematically. The fused from_logits = true path avoids the intermediate -y/p division and the separate softmax step. It therefore degrades more gracefully when a predicted probability is tiny. Do not mix the 2 pairings.

A softmax head with from_logits = true applies softmax twice. A linear head with from_logits = false feeds raw logits to a loss that expects probabilities. Loss details live in Loss Functions.

// More stable alternative to a Softmax head: emit logits, and fuse the softmax into the loss.
let mut model = builder
    .add(Dense::new(3, Activation::Linear).unwrap())      // raw logits, NOT probabilities
    .build(&Shape::known(x.shape()))?;
model.compile(optimizer, CategoricalCrossEntropy::new(true)); // from_logits = true

The activations are pure math, with no NaN/Inf sanitization. A NaN propagates through untouched. tanh saturates to 1 or -1 at large inputs. ReLU of a large negative input is 0. A NaN anywhere in a softmax row contaminates the whole row through the normalizer. Non-finite values surface downstream, as a NaN loss, not at the activation step.

3.2.5. PReLU: a learned negative slope

LeakyReLU fixes its negative slope at construction, so you have to guess a good value. PReLU learns it instead. The forward transform is the same: x for x >= 0 and alpha * x below 0. But alpha is a trainable array that the optimizer updates with the rest of the model.

PReLU is a layer, not an Activation variant, and it cannot become one. An Activation value carries no state, and this layer carries a trainable array. So it never fuses into Dense. Place it after the layer whose output it activates, which is the split form of 3.2.1.

PReLU::new(alpha) takes the starting slope alone. The build supplies the shape, exactly as it does for every other layer that owns an array. The layer then holds 1 slope per position of that shape with the batch axis removed, every one starting at alpha. A [batch, 32, 32, 64] build shape therefore gives 65,536 slopes, 1 per pixel and channel. That is almost never what you want after a convolution.

with_shared_axes(axes) is the fix. Each named axis drops to extent 1 in the slope array, and the slope broadcasts back over that axis. The axes count from the batch axis at 0, and the batch axis is shared already:

[1, 2] on a 4-D input is the standard choice after a convolution. It gives 1 slope per channel, which is the channel-wise form of He et al. (2015). [1, 2, 3] gives a single learned slope for the whole layer, which is LeakyReLU with the constant learned rather than guessed.

The layer also stops checking a shared axis at forward time, because 1 slope covers any extent. The same layer then serves images of several sizes. Every other axis after the batch axis must match the build shape, and a mismatch is Error::InvalidInput. The layer never checks the batch axis, so a partial final mini-batch always passes.

The derivative at exactly 0 is 0. That is neither branch: not the 1 the positive side gives, and not the alpha the negative side gives. It is what makes an alpha of 0 reproduce ReLU exactly, in the transform and in the gradient. That is why 0 is the natural starting value. LeakyReLU differs here, because it takes the positive branch at x >= 0 and its derivative at 0 is 1. A PReLU with a frozen uniform slope therefore matches LeakyReLU everywhere except at exactly 0.

Weight decay skips the slopes, in every optimizer that takes a weight_decay argument. Decay pulls a parameter toward 0. A slope of 0 turns the layer back into ReLU, so decay would erase what the layer learns. The slopes count as a no-decay parameter, the same class as a bias and a normalization gamma.

The stack below learns 1 slope per filter on a small convolutional model:

use ndarray::Array;
use rustyml::neural_network::Shape;
use rustyml::neural_network::layers::{Activation, Conv2D, Dense, Flatten, PReLU};
use rustyml::neural_network::losses::MeanSquaredError;
use rustyml::neural_network::optimizers::Adam;
use rustyml::neural_network::sequential::SequentialBuilder;

fn main() {
    // 2 single-channel 6x6 images, 1 continuous target each.
    let x = Array::from_shape_vec(
        (2, 6, 6, 1),
        (0..72).map(|v| 0.05 * v as f32 - 1.8).collect::<Vec<_>>(),
    )
    .unwrap()
    .into_dyn();
    let y = Array::from_shape_vec((2, 1), vec![0.3, -0.4])
        .unwrap()
        .into_dyn();

    let mut model = SequentialBuilder::new()
        // A 3x3 valid convolution over 6x6 emits [2, 4, 4, 4].
        .add(Conv2D::new(4, (3, 3), (1, 1), Activation::Linear).unwrap())
        // Sharing the 2 spatial axes gives 4 slopes, 1 per filter, instead of 64.
        .add(PReLU::new(0.25).unwrap().with_shared_axes(vec![1, 2]).unwrap())
        .add(Flatten::new())
        .add(Dense::new(1, Activation::Linear).unwrap())
        .build(&Shape::known(x.shape()))
        .unwrap();
    model.compile(
        Adam::new(0.01, 0.9, 0.999, 1e-8, 0.0).unwrap(),
        MeanSquaredError::new(),
    );

    model.fit(&x, &y, 10).unwrap();
    assert_eq!(model.predict(&x).unwrap().shape(), &[2, 1]);

    // The slope array keeps the extent-1 place of every shared axis.
    println!("alpha: {:?}", model.weight("1.alpha").unwrap().shape()); // [1, 1, 4]
}

The layer holds 1 array, named alpha. weight("alpha") reads it, and set_weights(alpha) writes it back. alpha is an ArrayD<f32>, and its rank follows the input rank, less 1 for the batch axis. A shared axis keeps its extent-1 place in that array, so a per-channel layer reports [1, 1, 64] and not [64].

3.2.6. Forward, backward, and the GEMM path

The forward pass is activation(input * W + b), computed as a single call into the gemmkit backend. The product, the per-column bias, and, for ReLU, the activation, all run in 1 pass. The bias and activation apply in the kernel’s epilogue while the output tile is still in registers. ReLU is the only one of the 12 activations with a fused epilogue. Linear needs no separate step, since the biased product is already the output. Every other activation runs as a separate vectorized Activation::forward pass over the biased product.

The fused result matches the unfused product plus scalar activation bit for bit, with 1 exception. The fused ReLU epilogue maps a NaN pre-activation to 0.0. The standalone ReLU activation used elsewhere in the crate instead propagates NaN. A NaN pre-activation already means a diverged model. The cached-output ReLU derivative then treats that 0.0 like a dead unit. It backpropagates a zero gradient there, instead of propagating the NaN further back.

Whether the product threads is gemmkit’s decision, not the layer’s. It compares batch * input_width * units against a work gate, 589,824 by default. That count is the raw product, not a FLOP count with a factor of 2. Below the gate, the product stays on 1 thread. Above it, the worker count ramps with the work, rather than moving straight to the full machine.

Small layers in a tight loop stay serial by design, because the per-call dispatch would dominate. These products also nest safely inside an outer parallel region. The backend is described in full in Matrix Multiplication, and the tunable gates in Performance Tuning and Parallelism.

The elementwise activation that follows has its own, separate parallelism gate. Each activation falls into 1 of 2 cost classes. ReLU, LeakyReLU, Softsign, and HardSigmoid are memory-bound “cheap maps” whose crossover sits at 4,000,000 elements. At any practical layer size, they run serial. Sigmoid, Tanh, Softmax, ELU, SELU, Softplus, and Exponential are exp-dominated, and go parallel above 131,072 elements. Linear copies the tensor and has no gate.

The backward pass uses the same 2 classes, but it picks the class from the derivative. The ELU, SELU, and Exponential derivatives are plain arithmetic on the cached output, with no exponential. Their backward pass therefore uses the cheap-map gate, even though their forward pass uses the exp gate. Moving either gate only trades serial for parallel. The results are identical, and every product is run-to-run deterministic on the same machine.

forward takes &self, and it parks the input and the activated output in the context. Ctx::training() selects the mode that parks them, and Ctx::inference() selects the mode that parks nothing. The 2 modes compute the same values, and the inference mode saves the cost of the parking. The layer itself holds no cache at all, which is why several threads can run inference through 1 model. The backward pass takes the parked output back and differentiates the activation with it.

It then computes 3 quantities. The weight gradient is input^T * grad, a GEMM. The bias gradient is the column-sum over the batch. The input gradient is grad * W^T, another GEMM. An input of rank 3 or more folds its leading axes into the row axis first. The 3 quantities therefore stay the same 3 products at every rank.

backward returns the input gradient, and it puts the other 2 in the gradient store of the context, under the names kernel and bias. Read one back with ctx.grads().get(ParamId::new(scope, name)), where scope is the position of the layer in the model that drives it. A bias-free layer computes no bias gradient, so the store holds none. The store sums, so a layer that 2 positions of a graph model share receives the total of the 2 backward passes.

Calling backward before forward returns Error::NeuralNetwork(NnError::ForwardPassNotRun). An upstream gradient whose shape does not match the parked output returns Error::ShapeMismatch. Both are errors, never panics.

3.2.7. Reading and setting weights

There is no standalone weight() / bias() getter of its own. The accessors come from the LayerBase trait, and every layer of the crate answers the same 3. All 3 need a built layer, because an unbuilt layer holds no array.

weights() gives a Vec<WeightRef<'_>>, 1 entry per array the layer holds. Each entry carries the name of the array, its kind, and a read view of the live storage. weights_mut() gives the same list for writing, under the same names and in the same order. weight(name) gives 1 read view, or None when the layer holds no array under that name. None of the 3 paths clones anything.

kind is a WeightKind, and it says whether an optimizer updates the array. Trainable covers a kernel, a bias, and a normalization scale or shift. NonTrainable covers state that the layer keeps and no optimizer writes. A Dense layer reports Trainable for both of its arrays.

parameters_mut() is the accessor that the optimizer takes, and the same trait holds it. It gives 1 ParamRef per trainable tensor: the name, a flat mutable view of the data, and a decays flag. The flag tells decoupled weight decay to skip a bias. The entry carries no gradient, because a gradient belongs to 1 pass and a parameter outlives every pass. The optimizer builds the address ParamId::new(scope, name), reads the gradient of that address out of the store, and skips a tensor that holds none.

kernel has shape (input_width, units), where the width is the one the build settled, and bias has shape (1, units), as the construction example in 3.2.2 shows. To write parameters, use set_weights(weights, bias), which validates both shapes. These same names are the on-disk addresses. A checkpoint writes each array as <scope>.<name>, where scope is the position of the layer. The kernel of the first layer of a model is therefore 0.kernel. Anything you can read here is what round-trips through save and load, in Saving and Loading Weights.

3.2.8. Two worked models

A regression net ends in a Linear head, and trains against mean squared error. The hidden layer fuses ReLU. The output layer fuses Linear, because a regression target is unbounded:

use ndarray::Array;
use rustyml::neural_network::Shape;
use rustyml::neural_network::layers::{Activation, Dense};
use rustyml::neural_network::losses::MeanSquaredError;
use rustyml::neural_network::optimizers::SGD;
use rustyml::neural_network::sequential::SequentialBuilder;

fn main() {
    // 4 samples, 3 features, 1 continuous target each.
    let x = Array::from_shape_vec(
        (4, 3),
        vec![0.0, 0.1, 0.2, 1.0, 0.9, 0.8, 0.2, 0.1, 0.0, 0.9, 1.0, 0.8],
    )
    .unwrap()
    .into_dyn();
    let y = Array::from_shape_vec((4, 1), vec![0.3, 2.7, 0.3, 2.7])
        .unwrap()
        .into_dyn();

    let mut model = SequentialBuilder::new()
        .add(Dense::new(8, Activation::ReLU).unwrap()) // hidden, fused ReLU
        .add(Dense::new(1, Activation::Linear).unwrap()) // regression head: identity
        .build(&Shape::known(x.shape()))
        .unwrap();
    model.compile(SGD::new(0.05, 0.9, false, 0.0).unwrap(), MeanSquaredError::new());

    model.fit(&x, &y, 20).unwrap();

    let preds = model.predict(&x).unwrap();
    println!("prediction shape: {:?}", preds.shape()); // [4, 1]
}

A classifier ends in a Softmax head over one-hot targets, paired with CategoricalCrossEntropy::new(false), because the head emits probabilities:

use ndarray::Array;
use rustyml::neural_network::Shape;
use rustyml::neural_network::layers::{Activation, Dense};
use rustyml::neural_network::losses::CategoricalCrossEntropy;
use rustyml::neural_network::optimizers::Adam;
use rustyml::neural_network::sequential::SequentialBuilder;

fn main() {
    // 4 samples, 3 features, 2 classes (one-hot targets).
    let x = Array::from_shape_vec(
        (4, 3),
        vec![0.0, 0.1, 0.2, 1.0, 0.9, 0.8, 0.1, 0.0, 0.2, 0.8, 1.0, 0.9],
    )
    .unwrap()
    .into_dyn();
    let y = Array::from_shape_vec((4, 2), vec![1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0])
        .unwrap()
        .into_dyn();

    let mut model = SequentialBuilder::new()
        .add(Dense::new(8, Activation::ReLU).unwrap())
        .add(Dense::new(2, Activation::Softmax { axis: -1 }).unwrap()) // probability head
        .build(&Shape::known(x.shape()))
        .unwrap();
    model.compile(
        Adam::new(0.01, 0.9, 0.999, 1e-8, 0.0).unwrap(),
        CategoricalCrossEntropy::new(false),
    );

    model.fit(&x, &y, 20).unwrap();

    let probs = model.predict(&x).unwrap();
    // Each row is a distribution over the 2 classes, summing to 1.
    println!("class-probability shape: {:?}", probs.shape()); // [4, 2]
}

Swap the head to Activation::Linear, and the loss to CategoricalCrossEntropy::new(true). This classifier then trains on the more stable fused-logits path, and predict produces raw scores instead of probabilities. Which optimizer to compile with, and how learning rate and momentum interact with the dead-neuron and saturation behavior above, is the subject of Optimizers.