rustyml 0.15.0

A high-performance machine learning & deep learning library in pure Rust, offering ML algorithms and neural network support
Documentation
# 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](./3.1._The_Sequential_Model.md). The layers here are what you stack inside a `Sequential` model.

## 3.2.1. The fused-activation design

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

```rust,ignore
pub fn new(
    input_dim: usize,
    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 one struct (see [Saving and Loading Weights](./3.9._Saving_and_Loading_Weights.md)). 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](./3.8._Regularization_and_Normalization_Layers.md)). The 2 forms are equivalent:

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

model.add(Dense::new(64, 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(input_dim, units, activation)` returns `Result<Dense, Error>`. Both dimensions must be non-zero. Passing `0` for either yields `Error::InvalidParameter` (see [Error Handling](../Chapter-01/1.6._Error_Handling.md)). `input_dim` is the number of features per row. `units` is the number of neurons, and therefore the output width.

The parameter count is `input_dim * units + units`: one weight per (input, output) pair, plus one bias per output. `param_count()` reports it as `TrainingParameters::Trainable(n)`. `output_shape()` renders `(None, units)`, where `None` is the dynamic batch dimension, mirroring Keras' summary. A `Dense(4, 3, ...)` therefore holds a `4 x 3` weight matrix and a `1 x 3` bias. That is `12 + 3 = 15` trainable scalars.

Input must be a 2-D tensor of shape `(batch, input_dim)`. A 1-D or 3-D input is rejected with `Error::InvalidInput`, rather than being silently reshaped. If you are feeding the output of a convolutional or recurrent stack, flatten it to 2 dimensions first. The following program constructs a layer, inspects it, and injects known weights:

```rust
use ndarray::Array2;
use rustyml::neural_network::layers::layer_weight::LayerWeight;
use rustyml::neural_network::layers::{Activation, Dense, TrainingParameters};
use rustyml::neural_network::traits::Layer;

fn main() {
    // A 4 -> 3 dense layer with a fused ReLU activation.
    let mut dense = Dense::new(4, 3, Activation::ReLU).unwrap();

    // 4*3 weights + 3 bias = 15 trainable scalars.
    assert_eq!(dense.param_count(), TrainingParameters::Trainable(15));
    println!("output shape: {}", dense.output_shape()); // (None, 3)

    // Read the freshly initialized parameters without cloning them.
    match dense.get_weights() {
        LayerWeight::Dense(w) => {
            println!("weight {:?}, bias {:?}", w.weight.shape(), w.bias.shape());
        }
        _ => unreachable!(),
    }

    // Inject known weights. The shapes are validated against the layer's config.
    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();
}
```

`set_weights` checks both shapes and returns `Error::NeuralNetwork(NnError::WeightShape)` on a mismatch. A `(3, 3)` weight for a `4 -> 3` layer, or a `(1, 4)` bias, is rejected rather than truncated.

## 3.2.3. Weight initialization

Weights use **Xavier/Glorot uniform** initialization. Each element is drawn from `Uniform(-limit, +limit)`, where `limit = sqrt(6 / (input_dim + units))`. Biases start at exactly zero. This is the fan-in-plus-fan-out Glorot scheme. It matches Keras' `Dense` default exactly (`glorot_uniform` kernel, `zeros` bias).

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. By default the seed comes from the process-global seed if one is set, otherwise from entropy. Two 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 construct the model. This also fixes dropout masks and the fit-time batch shuffle.

You can instead seed one layer explicitly, with `Dense::new(...)?.with_random_state(seed)`. This re-runs Glorot with that seed and leaves the global stream untouched. The full seeding model is in [Reproducibility and Random Seeds](../Chapter-07/7.1._Reproducibility_and_Random_Seeds.md). 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.

| `Activation` | Forward `f(x)` | Backward (given upstream `g`, output `a`) | Output range |
| --- | --- | --- | --- |
| `ReLU` | `max(0, x)` | pass `g` where `a > 0`, else `0` | `[0, inf)` |
| `Sigmoid` | `1 / (1 + e^-x)` | `g * a * (1 - a)` | `(0, 1)` |
| `Tanh` | `tanh(x)` | `g * (1 - a^2)` | `(-1, 1)` |
| `Softmax` | shifted `exp`, row-normalized | `a_i * (g_i - sum_j(a_j * g_j))` (row Jacobian) | simplex, each row sums to `1` |
| `Linear` | `x` | pass `g` through | `(-inf, inf)` |
| `LeakyReLU { negative_slope }` | `x` for `x >= 0`, else `negative_slope * x` | pass `g` where `a >= 0`, else `g * negative_slope` | `(-inf, inf)` |
| `ELU { alpha }` | `x` for `x > 0`, else `alpha * (e^x - 1)` | pass `g` where `a > 0`, else `g * (a + alpha)` | `(-alpha, inf)` |
| `SELU` | `scale * x` for `x > 0`, else `scale * alpha * (e^x - 1)` | `g * scale` where `a > 0`, else `g * (a + scale * alpha)` | `(-scale * alpha, inf)` |
| `Softplus` | `ln(1 + e^x)` | `g * (1 - e^-a)` | `(0, inf)` |
| `Softsign` | `x / (1 + abs(x))` | `g * (1 - abs(a))^2` | `(-1, 1)` |
| `HardSigmoid` | `clip(x/6 + 0.5, 0, 1)` | `g / 6` where `0 < a < 1`, else `0` | `[0, 1]` |
| `Exponential` | `e^x` | `g * a` | `(0, inf)` |

**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:

```rust
use ndarray::Array;
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::Sequential;

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 = Sequential::new();
    model
        // Fused: the activation runs inside the Dense layer.
        .add(Dense::new(3, 8, Activation::ELU { alpha: 1.0 }).unwrap())
        // Split: an identity Dense, then the same activation as its own layer.
        .add(Dense::new(8, 8, Activation::Linear).unwrap())
        .add(ELU::new(1.0).unwrap())
        .add(Dense::new(8, 1, Activation::Linear).unwrap()) // regression head
        .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(3, 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 row of logits into a probability distribution over the last axis. The forward pass subtracts each row'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 needs at least a 2-D input. A 1-D tensor returns `Error::InvalidInput`. Its backward pass is the true Jacobian-vector product across the row, not an elementwise multiply, and each gradient row sums to zero.

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 two 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](./3.3._Loss_Functions.md).

```rust,ignore
// More stable alternative to a Softmax head: emit logits, and fuse the softmax into the loss.
model
    .add(Dense::new(8, 3, Activation::Linear).unwrap())    // raw logits, NOT probabilities
    .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(input_shape, alpha)` takes the input shape with the batch axis first, exactly as the convolutional layers do. It then holds 1 slope per position of that shape with the batch axis removed, every one starting at `alpha`. A `[batch, 32, 32, 64]` input 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:

| `input_shape` | `shared_axes` | Slope shape | Slope count |
| --- | --- | --- | --- |
| `[batch, 8]` | none | `[8]` | 8 |
| `[batch, 32, 32, 64]` | none | `[32, 32, 64]` | 65,536 |
| `[batch, 32, 32, 64]` | `[1, 2]` | `[1, 1, 64]` | 64 |
| `[batch, 32, 32, 64]` | `[1, 2, 3]` | `[1, 1, 1]` | 1 |
| `[batch, 20, 16]` | `[1]` | `[1, 16]` | 16 |

`[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.

A shared axis also stops being checked 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 `input_shape`, and a mismatch is `Error::InvalidInput`. The batch axis is never checked, 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:

```rust
use ndarray::Array;
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::Sequential;

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 = Sequential::new();
    model
        // A 3x3 valid convolution over 6x6 emits [2, 4, 4, 4].
        .add(Conv2D::new(4, (3, 3), vec![2, 6, 6, 1], (1, 1), Activation::Linear).unwrap())
        // Sharing the 2 spatial axes gives 4 slopes, 1 per filter, instead of 64.
        .add(
            PReLU::new(vec![2, 4, 4, 4], 0.25)
                .unwrap()
                .with_shared_axes(vec![1, 2])
                .unwrap(),
        )
        .add(Flatten::new(vec![2, 4, 4, 4]).unwrap())
        .add(Dense::new(64, 1, Activation::Linear).unwrap())
        .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]);
}
```

`get_weights()` returns `LayerWeight::PReLU(PReLULayerWeight { alpha })`, 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 returns `[1, 1, 64]` and not `[64]`.

| Error | Cause |
| --- | --- |
| `Error::InvalidInput` | `input_shape` has rank below 2, or holds a `0` |
| `Error::InvalidParameter` | `alpha` is not finite |
| `Error::InvalidParameter` | `shared_axes` holds `0`, an axis at or above the rank, or a repeat |
| `Error::InvalidInput` | a forward input whose rank differs, or whose extent on an axis that is not shared differs, from `input_shape` |
| `Error::NeuralNetwork(NnError::WeightShape)` | `set_weights` got an array that is not the slope shape |

## 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](https://docs.rs/gemmkit) backend. The product, the per-column bias, and, for `ReLU`, the activation, all run in one 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 one 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_dim * 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 one 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](../Chapter-06/6.2._Matrix_Multiplication.md), and the tunable gates in [Performance Tuning and Parallelism](../Chapter-07/7.3._Performance_Tuning_and_Parallelism.md).

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` caches the input and the activated output for the backward pass. `predict` is the eval-mode twin. It computes the same values, but writes no caches. The backward pass first differentiates the activation using the cached output.

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.

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

## 3.2.7. Reading and setting weights

There is no standalone `weights()` / `bias()` getter. The accessor is `get_weights()` from the `Layer` trait. It returns a `LayerWeight` enum. For a dense layer, that is `LayerWeight::Dense(DenseLayerWeight { weight, bias })`, where `weight` and `bias` are `Cow<Array2<f32>>` borrowed from the live layer, with no clone.

`weight` has shape `(input_dim, units)`, 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. This same `LayerWeight` enum is the on-disk weight format. Anything you can read here is what round-trips through save and load, in [Saving and Loading Weights](./3.9._Saving_and_Loading_Weights.md).

## 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:

```rust
use ndarray::Array;
use rustyml::neural_network::layers::{Activation, Dense};
use rustyml::neural_network::losses::MeanSquaredError;
use rustyml::neural_network::optimizers::SGD;
use rustyml::neural_network::sequential::Sequential;

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 = Sequential::new();
    model
        .add(Dense::new(3, 8, Activation::ReLU).unwrap()) // hidden, fused ReLU
        .add(Dense::new(8, 1, Activation::Linear).unwrap()) // regression head: identity
        .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:

```rust
use ndarray::Array;
use rustyml::neural_network::layers::{Activation, Dense};
use rustyml::neural_network::losses::CategoricalCrossEntropy;
use rustyml::neural_network::optimizers::Adam;
use rustyml::neural_network::sequential::Sequential;

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 = Sequential::new();
    model
        .add(Dense::new(3, 8, Activation::ReLU).unwrap())
        .add(Dense::new(8, 2, Activation::Softmax).unwrap()) // probability head
        .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](./3.4._Optimizers.md).