Skip to main content

ipfrs_tensorlogic/
feed_forward.rs

1//! Feedforward network layer for transformer blocks.
2//!
3//! This module implements the position-wise feedforward sub-layer that appears
4//! in every transformer encoder/decoder block (Vaswani et al. 2017).  The
5//! standard two-layer structure is:
6//!
7//! ```text
8//! output = Linear2( Activation( Linear1( input ) ) )
9//! ```
10//!
11//! where `Linear1` projects from `input_dim` → `hidden_dim` (typically
12//! `4 × input_dim`) and `Linear2` projects back to `output_dim`.
13//!
14//! # Supported activations
15//!
16//! | Variant   | Description                                    |
17//! |-----------|------------------------------------------------|
18//! | `ReLU`    | max(0, x)                                      |
19//! | `GELU`    | Gaussian Error Linear Unit (tanh approximation)|
20//! | `SiLU`    | x · σ(x)  (also known as Swish)               |
21//! | `Linear`  | Identity — no activation applied               |
22//!
23//! # Weight initialisation
24//!
25//! Weights are initialised with **He (Kaiming) normal** initialisation using a
26//! custom xorshift64 PRNG with Box-Muller transform — zero external crate
27//! dependency.
28//!
29//! # Example
30//!
31//! ```rust
32//! use ipfrs_tensorlogic::{FeedForwardConfig, FeedForwardActivation, FeedForwardNetwork};
33//!
34//! let cfg = FeedForwardConfig {
35//!     input_dim: 8,
36//!     hidden_dim: 32,
37//!     output_dim: 8,
38//!     activation: FeedForwardActivation::GELU,
39//!     use_bias: true,
40//!     dropout_rate: 0.1,
41//! };
42//!
43//! let mut net = FeedForwardNetwork::new(cfg, 42);
44//!
45//! let token = vec![1.0_f64; 8];
46//! let out = net.forward(&token);
47//! assert_eq!(out.len(), 8);
48//! ```
49
50use std::f64::consts::PI;
51
52// ── Activation enum ───────────────────────────────────────────────────────────
53
54/// Activation function applied between the two linear projections.
55#[derive(Debug, Clone, PartialEq)]
56pub enum FeedForwardActivation {
57    /// Rectified Linear Unit: max(0, x).
58    ReLU,
59    /// Gaussian Error Linear Unit (tanh approximation).
60    GELU,
61    /// Sigmoid Linear Unit / Swish: x · σ(x).
62    SiLU,
63    /// No activation — passes the pre-activation values through unchanged.
64    Linear,
65}
66
67// ── Configuration ─────────────────────────────────────────────────────────────
68
69/// Configuration for [`FeedForwardNetwork`].
70#[derive(Debug, Clone)]
71pub struct FeedForwardConfig {
72    /// Dimensionality of the input vector (and typically the output as well).
73    pub input_dim: usize,
74    /// Width of the hidden (intermediate) projection; commonly `4 × input_dim`.
75    pub hidden_dim: usize,
76    /// Dimensionality of the network output.
77    pub output_dim: usize,
78    /// Activation function applied after the first linear transformation.
79    pub activation: FeedForwardActivation,
80    /// When `true` bias vectors are allocated and applied in each layer.
81    pub use_bias: bool,
82    /// Conceptual dropout rate; stored for reference — no stochastic drop is
83    /// applied during deterministic inference or the current forward pass.
84    pub dropout_rate: f64,
85}
86
87// ── Single layer ──────────────────────────────────────────────────────────────
88
89/// A single affine (linear) layer: weight matrix and optional bias.
90///
91/// `weights` has shape `out_dim × in_dim` — each row is the weight vector
92/// for one output neuron.
93#[derive(Debug, Clone)]
94pub struct FFLayer {
95    /// Weight matrix stored row-major: `weights[o][i]` = weight from input `i`
96    /// to output neuron `o`.  Shape: `out_dim × in_dim`.
97    pub weights: Vec<Vec<f64>>,
98    /// Bias vector of length `out_dim`.  All zeros when `use_bias` is `false`.
99    pub bias: Vec<f64>,
100}
101
102// ── Running statistics ────────────────────────────────────────────────────────
103
104/// Lightweight counters accumulated across [`FeedForwardNetwork::forward`] calls.
105#[derive(Debug, Clone, Default)]
106pub struct FFStats {
107    /// Total number of times `forward` or `forward_batch` has been invoked.
108    pub total_forward_calls: u64,
109    /// Cumulative number of individual tokens (vectors) that have passed
110    /// through the network.
111    pub total_tokens_processed: u64,
112}
113
114// ── Network ───────────────────────────────────────────────────────────────────
115
116/// Two-layer feedforward network suitable for use as the FFN sub-layer of a
117/// transformer block.
118pub struct FeedForwardNetwork {
119    config: FeedForwardConfig,
120    /// First projection: `input_dim` → `hidden_dim`.
121    layer1: FFLayer,
122    /// Second projection: `hidden_dim` → `output_dim`.
123    layer2: FFLayer,
124    /// Internal xorshift64 PRNG state (retained for reproducible re-init).
125    rng_state: u64,
126    stats: FFStats,
127}
128
129// ── Implementation ────────────────────────────────────────────────────────────
130
131impl FeedForwardNetwork {
132    /// Construct a new network from `config`, initialising weights with He
133    /// normal initialisation seeded by `seed`.
134    pub fn new(config: FeedForwardConfig, seed: u64) -> Self {
135        let mut rng = if seed == 0 { 0x853c49e6748fea9b } else { seed };
136
137        let layer1 = Self::init_layer(
138            config.input_dim,
139            config.hidden_dim,
140            config.use_bias,
141            &mut rng,
142        );
143        let layer2 = Self::init_layer(
144            config.hidden_dim,
145            config.output_dim,
146            config.use_bias,
147            &mut rng,
148        );
149
150        Self {
151            config,
152            layer1,
153            layer2,
154            rng_state: rng,
155            stats: FFStats::default(),
156        }
157    }
158
159    /// Run a single token (flat `f64` slice of length `input_dim`) through the
160    /// network and return the output vector of length `output_dim`.
161    ///
162    /// If the input length does not match `input_dim` the network applies
163    /// whatever it can and pads/truncates gracefully — no panic.
164    pub fn forward(&mut self, input: &[f64]) -> Vec<f64> {
165        self.stats.total_forward_calls += 1;
166        self.stats.total_tokens_processed += 1;
167
168        // Layer 1: input_dim → hidden_dim
169        let mut hidden = Self::linear_transform(input, &self.layer1);
170
171        // Apply activation element-wise
172        for v in hidden.iter_mut() {
173            *v = self.apply_activation(*v);
174        }
175
176        // Layer 2: hidden_dim → output_dim
177        Self::linear_transform(&hidden, &self.layer2)
178    }
179
180    /// Run a batch of tokens through the network.
181    ///
182    /// Returns one output vector per input token; empty input yields an empty
183    /// result with no panic.
184    pub fn forward_batch(&mut self, inputs: &[Vec<f64>]) -> Vec<Vec<f64>> {
185        self.stats.total_forward_calls += 1;
186        let token_count = inputs.len() as u64;
187        self.stats.total_tokens_processed += token_count;
188
189        inputs
190            .iter()
191            .map(|token| {
192                // Layer 1
193                let mut hidden = Self::linear_transform(token, &self.layer1);
194                for v in hidden.iter_mut() {
195                    *v = self.apply_activation(*v);
196                }
197                // Layer 2
198                Self::linear_transform(&hidden, &self.layer2)
199            })
200            .collect()
201    }
202
203    /// Affine transformation: `output[o] = bias[o] + Σ_i weights[o][i] * input[i]`.
204    ///
205    /// Dimension mismatches are handled gracefully: the dot-product iterates
206    /// over `min(in_dim, input.len())` elements and missing bias values default
207    /// to `0.0`.
208    pub fn linear_transform(input: &[f64], layer: &FFLayer) -> Vec<f64> {
209        layer
210            .weights
211            .iter()
212            .enumerate()
213            .map(|(o, row)| {
214                let dot: f64 = row.iter().zip(input.iter()).map(|(w, x)| w * x).sum();
215                let b = layer.bias.get(o).copied().unwrap_or(0.0);
216                dot + b
217            })
218            .collect()
219    }
220
221    /// Apply the configured activation function to a single scalar.
222    #[inline]
223    pub fn apply_activation(&self, x: f64) -> f64 {
224        match self.config.activation {
225            FeedForwardActivation::ReLU => Self::relu(x),
226            FeedForwardActivation::GELU => Self::gelu(x),
227            FeedForwardActivation::SiLU => Self::silu(x),
228            FeedForwardActivation::Linear => x,
229        }
230    }
231
232    /// Rectified Linear Unit.
233    #[inline]
234    pub fn relu(x: f64) -> f64 {
235        x.max(0.0)
236    }
237
238    /// GELU using the tanh approximation (Hendrycks & Gimpel 2016):
239    ///
240    /// ```text
241    /// GELU(x) ≈ 0.5 · x · (1 + tanh(√(2/π) · (x + 0.044715 · x³)))
242    /// ```
243    pub fn gelu(x: f64) -> f64 {
244        let c = (2.0_f64 / PI).sqrt();
245        let inner = c * (x + 0.044715 * x * x * x);
246        0.5 * x * (1.0 + inner.tanh())
247    }
248
249    /// Sigmoid Linear Unit (Swish): `x · σ(x)`.
250    #[inline]
251    pub fn silu(x: f64) -> f64 {
252        x * Self::sigmoid(x)
253    }
254
255    /// Logistic sigmoid: `σ(x) = 1 / (1 + e^{−x})`.
256    #[inline]
257    pub fn sigmoid(x: f64) -> f64 {
258        1.0 / (1.0 + (-x).exp())
259    }
260
261    /// Initialise a single [`FFLayer`] with He (Kaiming) normal weights.
262    ///
263    /// He init draws weights from N(0, σ²) where σ = √(2 / in_dim).
264    /// Bias is always zero-initialised.
265    pub fn init_layer(in_dim: usize, out_dim: usize, use_bias: bool, rng: &mut u64) -> FFLayer {
266        // He-init standard deviation: sqrt(2 / fan_in)
267        let std_dev = if in_dim > 0 {
268            (2.0_f64 / in_dim as f64).sqrt()
269        } else {
270            1.0
271        };
272
273        let weights: Vec<Vec<f64>> = (0..out_dim)
274            .map(|_| {
275                (0..in_dim)
276                    .map(|_| Self::next_normal(rng) * std_dev)
277                    .collect()
278            })
279            .collect();
280
281        let bias = if use_bias {
282            vec![0.0_f64; out_dim]
283        } else {
284            // Even when use_bias is false we keep a zero vector so that
285            // `linear_transform` never needs to branch on this field.
286            vec![0.0_f64; out_dim]
287        };
288
289        FFLayer { weights, bias }
290    }
291
292    /// Draw a standard-normal sample using the xorshift64 PRNG (Marsaglia 2003)
293    /// and Box-Muller transform.
294    ///
295    /// Two uniform samples `u1`, `u2` ∈ (0, 1] are generated; one standard-
296    /// normal deviate is returned.
297    pub fn next_normal(rng: &mut u64) -> f64 {
298        let u1 = Self::xorshift64(rng);
299        let u2 = Self::xorshift64(rng);
300
301        // Box-Muller: Z = √(-2 ln u1) · cos(2π u2)
302        let r = (-2.0 * u1.ln()).sqrt();
303        r * (2.0 * PI * u2).cos()
304    }
305
306    /// xorshift64 PRNG step (Marsaglia 2003) — returns a uniform f64 in (0, 1].
307    #[inline]
308    fn xorshift64(state: &mut u64) -> f64 {
309        let mut x = *state;
310        // Guard against zero state (would produce all-zero sequence)
311        if x == 0 {
312            x = 0x853c49e6748fea9b;
313        }
314        x ^= x << 13;
315        x ^= x >> 7;
316        x ^= x << 17;
317        *state = x;
318
319        // Map to (0, 1] — divide by 2^64, add tiny epsilon to exclude 0
320        (x as f64) / (u64::MAX as f64) + f64::EPSILON
321    }
322
323    /// Reinitialise both layers using the stored `rng_state`, effectively
324    /// resetting weights to a new He-normal draw that continues from where the
325    /// original seed sequence left off.  Useful for experimentation without
326    /// constructing a brand-new network.
327    pub fn reinit_weights(&mut self) {
328        let mut rng = self.rng_state;
329        self.layer1 = Self::init_layer(
330            self.config.input_dim,
331            self.config.hidden_dim,
332            self.config.use_bias,
333            &mut rng,
334        );
335        self.layer2 = Self::init_layer(
336            self.config.hidden_dim,
337            self.config.output_dim,
338            self.config.use_bias,
339            &mut rng,
340        );
341        self.rng_state = rng;
342    }
343
344    /// Reference to the accumulated runtime statistics.
345    pub fn stats(&self) -> &FFStats {
346        &self.stats
347    }
348}
349
350// ── Tests ─────────────────────────────────────────────────────────────────────
351
352#[cfg(test)]
353mod tests {
354    use super::*;
355
356    // ── Helpers ───────────────────────────────────────────────────────────────
357
358    fn make_net(
359        in_dim: usize,
360        hidden: usize,
361        out_dim: usize,
362        act: FeedForwardActivation,
363    ) -> FeedForwardNetwork {
364        let cfg = FeedForwardConfig {
365            input_dim: in_dim,
366            hidden_dim: hidden,
367            output_dim: out_dim,
368            activation: act,
369            use_bias: true,
370            dropout_rate: 0.0,
371        };
372        FeedForwardNetwork::new(cfg, 12345)
373    }
374
375    fn make_net_no_bias(in_dim: usize, hidden: usize, out_dim: usize) -> FeedForwardNetwork {
376        let cfg = FeedForwardConfig {
377            input_dim: in_dim,
378            hidden_dim: hidden,
379            output_dim: out_dim,
380            activation: FeedForwardActivation::Linear,
381            use_bias: false,
382            dropout_rate: 0.0,
383        };
384        FeedForwardNetwork::new(cfg, 99)
385    }
386
387    // ── 1. Output shape — single token ────────────────────────────────────────
388
389    #[test]
390    fn test_forward_output_shape() {
391        let mut net = make_net(8, 32, 8, FeedForwardActivation::ReLU);
392        let out = net.forward(&[1.0; 8]);
393        assert_eq!(out.len(), 8);
394    }
395
396    // ── 2. Output shape — batch ───────────────────────────────────────────────
397
398    #[test]
399    fn test_forward_batch_shape() {
400        let mut net = make_net(4, 16, 4, FeedForwardActivation::GELU);
401        let batch: Vec<Vec<f64>> = (0..5).map(|_| vec![1.0; 4]).collect();
402        let out = net.forward_batch(&batch);
403        assert_eq!(out.len(), 5);
404        for row in &out {
405            assert_eq!(row.len(), 4);
406        }
407    }
408
409    // ── 3. Linear transform — known values ───────────────────────────────────
410
411    #[test]
412    fn test_linear_transform_known_values() {
413        // weights = [[1, 0], [0, 1]], bias = [10, 20]
414        let layer = FFLayer {
415            weights: vec![vec![1.0, 0.0], vec![0.0, 1.0]],
416            bias: vec![10.0, 20.0],
417        };
418        let input = vec![3.0, 7.0];
419        let out = FeedForwardNetwork::linear_transform(&input, &layer);
420        assert!((out[0] - 13.0).abs() < 1e-12, "expected 13, got {}", out[0]);
421        assert!((out[1] - 27.0).abs() < 1e-12, "expected 27, got {}", out[1]);
422    }
423
424    // ── 4. Linear transform — scaling ────────────────────────────────────────
425
426    #[test]
427    fn test_linear_transform_scaling() {
428        let layer = FFLayer {
429            weights: vec![vec![2.0, 3.0]],
430            bias: vec![0.0],
431        };
432        let input = vec![4.0, 5.0];
433        let out = FeedForwardNetwork::linear_transform(&input, &layer);
434        assert!((out[0] - 23.0).abs() < 1e-12);
435    }
436
437    // ── 5. ReLU: positive ────────────────────────────────────────────────────
438
439    #[test]
440    fn test_relu_positive() {
441        assert!((FeedForwardNetwork::relu(3.5) - 3.5).abs() < 1e-12);
442    }
443
444    // ── 6. ReLU: negative ────────────────────────────────────────────────────
445
446    #[test]
447    fn test_relu_negative() {
448        assert!((FeedForwardNetwork::relu(-2.0)).abs() < 1e-12);
449    }
450
451    // ── 7. ReLU: zero ────────────────────────────────────────────────────────
452
453    #[test]
454    fn test_relu_zero() {
455        assert!((FeedForwardNetwork::relu(0.0)).abs() < 1e-12);
456    }
457
458    // ── 8. GELU: zero ────────────────────────────────────────────────────────
459
460    #[test]
461    fn test_gelu_zero() {
462        // GELU(0) = 0
463        assert!(FeedForwardNetwork::gelu(0.0).abs() < 1e-10);
464    }
465
466    // ── 9. GELU: large positive ───────────────────────────────────────────────
467
468    #[test]
469    fn test_gelu_large_positive() {
470        // For large x, GELU(x) ≈ x
471        let x = 10.0_f64;
472        let g = FeedForwardNetwork::gelu(x);
473        assert!((g - x).abs() < 1e-4, "GELU({}) = {} expected ≈ {}", x, g, x);
474    }
475
476    // ── 10. GELU: large negative ──────────────────────────────────────────────
477
478    #[test]
479    fn test_gelu_large_negative() {
480        // For large negative x, GELU(x) ≈ 0
481        let g = FeedForwardNetwork::gelu(-10.0);
482        assert!(g.abs() < 1e-4, "GELU(-10) = {} expected ≈ 0", g);
483    }
484
485    // ── 11. SiLU: zero ───────────────────────────────────────────────────────
486
487    #[test]
488    fn test_silu_zero() {
489        // SiLU(0) = 0 · σ(0) = 0 · 0.5 = 0
490        assert!(FeedForwardNetwork::silu(0.0).abs() < 1e-12);
491    }
492
493    // ── 12. SiLU: positive ───────────────────────────────────────────────────
494
495    #[test]
496    fn test_silu_positive() {
497        // SiLU(1) = 1 · σ(1) ≈ 0.7310585786300049
498        let s = FeedForwardNetwork::silu(1.0);
499        assert!((s - 0.7310585786300049).abs() < 1e-9);
500    }
501
502    // ── 13. SiLU: large positive ≈ identity ──────────────────────────────────
503
504    #[test]
505    fn test_silu_large_positive() {
506        // SiLU(x) → x as x → +∞
507        let x = 20.0_f64;
508        let s = FeedForwardNetwork::silu(x);
509        assert!((s - x).abs() < 1e-4);
510    }
511
512    // ── 14. Linear activation ─────────────────────────────────────────────────
513
514    #[test]
515    fn test_linear_activation_identity() {
516        let net = make_net(4, 8, 4, FeedForwardActivation::Linear);
517        // apply_activation should be identity for Linear
518        assert!((net.apply_activation(3.7) - 3.7).abs() < 1e-12);
519        assert!((net.apply_activation(-1.23) - (-1.23)).abs() < 1e-12);
520    }
521
522    // ── 15. Bias addition ─────────────────────────────────────────────────────
523
524    #[test]
525    fn test_bias_addition() {
526        let layer = FFLayer {
527            weights: vec![vec![0.0, 0.0], vec![0.0, 0.0]],
528            bias: vec![5.0, -3.0],
529        };
530        let input = vec![1.0, 2.0];
531        let out = FeedForwardNetwork::linear_transform(&input, &layer);
532        assert!((out[0] - 5.0).abs() < 1e-12);
533        assert!((out[1] - (-3.0)).abs() < 1e-12);
534    }
535
536    // ── 16. He init — weights non-zero ───────────────────────────────────────
537
538    #[test]
539    fn test_he_init_weights_nonzero() {
540        let mut rng = 42_u64;
541        let layer = FeedForwardNetwork::init_layer(8, 16, true, &mut rng);
542        let nonzero = layer.weights.iter().flatten().any(|&w| w.abs() > 1e-12);
543        assert!(nonzero, "All weights were zero — He init failed");
544    }
545
546    // ── 17. He init — bias is zero ────────────────────────────────────────────
547
548    #[test]
549    fn test_he_init_bias_zero() {
550        let mut rng = 42_u64;
551        let layer = FeedForwardNetwork::init_layer(8, 16, true, &mut rng);
552        for b in &layer.bias {
553            assert!(b.abs() < 1e-30, "Bias should be zero-initialised");
554        }
555    }
556
557    // ── 18. He init — weight variance ─────────────────────────────────────────
558
559    #[test]
560    fn test_he_init_variance_property() {
561        // Variance of He-normal weights ≈ 2 / fan_in
562        let fan_in = 64_usize;
563        let fan_out = 256_usize;
564        let mut rng = 7654321_u64;
565        let layer = FeedForwardNetwork::init_layer(fan_in, fan_out, true, &mut rng);
566
567        let all: Vec<f64> = layer.weights.into_iter().flatten().collect();
568        let n = all.len() as f64;
569        let mean = all.iter().sum::<f64>() / n;
570        let variance = all.iter().map(|w| (w - mean).powi(2)).sum::<f64>() / n;
571        let expected_var = 2.0 / fan_in as f64;
572
573        // Allow 50 % relative tolerance — empirical sampling noise
574        assert!(
575            (variance - expected_var).abs() / expected_var < 0.5,
576            "Variance {:.4} too far from expected {:.4}",
577            variance,
578            expected_var
579        );
580    }
581
582    // ── 19. Single token forward — output is finite ───────────────────────────
583
584    #[test]
585    fn test_single_token_forward_finite() {
586        let mut net = make_net(16, 64, 16, FeedForwardActivation::GELU);
587        let token: Vec<f64> = (0..16).map(|i| i as f64 * 0.1).collect();
588        let out = net.forward(&token);
589        for (i, v) in out.iter().enumerate() {
590            assert!(v.is_finite(), "output[{}] = {} is not finite", i, v);
591        }
592    }
593
594    // ── 20. Sequential batch — each result matches individual forward ─────────
595
596    #[test]
597    fn test_sequential_batch_matches_individual() {
598        // Use a fresh net for each mode to ensure identical RNG state influence.
599        // Because forward() updates stats but not weights, results must be equal.
600        let mut net = make_net(4, 8, 4, FeedForwardActivation::SiLU);
601        let tokens: Vec<Vec<f64>> = vec![
602            vec![1.0, 0.0, -1.0, 0.5],
603            vec![0.0, 1.0, 0.0, -1.0],
604            vec![0.5, 0.5, 0.5, 0.5],
605        ];
606
607        let batch_out = net.forward_batch(&tokens);
608
609        // Reset stats to compare only outputs
610        let mut net2 = make_net(4, 8, 4, FeedForwardActivation::SiLU);
611        let individual: Vec<Vec<f64>> = tokens.iter().map(|t| net2.forward(t)).collect();
612
613        for (b, ind) in batch_out.iter().zip(individual.iter()) {
614            for (bv, iv) in b.iter().zip(ind.iter()) {
615                assert!(
616                    (bv - iv).abs() < 1e-12,
617                    "batch vs individual mismatch: {} vs {}",
618                    bv,
619                    iv
620                );
621            }
622        }
623    }
624
625    // ── 21. Stats tracking — forward_calls ───────────────────────────────────
626
627    #[test]
628    fn test_stats_forward_calls() {
629        let mut net = make_net(4, 8, 4, FeedForwardActivation::ReLU);
630        assert_eq!(net.stats().total_forward_calls, 0);
631        net.forward(&[0.0; 4]);
632        net.forward(&[1.0; 4]);
633        assert_eq!(net.stats().total_forward_calls, 2);
634    }
635
636    // ── 22. Stats tracking — tokens processed ────────────────────────────────
637
638    #[test]
639    fn test_stats_tokens_processed() {
640        let mut net = make_net(4, 8, 4, FeedForwardActivation::ReLU);
641        let batch: Vec<Vec<f64>> = (0..7).map(|_| vec![0.0; 4]).collect();
642        net.forward_batch(&batch);
643        // forward_batch counts all tokens
644        assert_eq!(net.stats().total_tokens_processed, 7);
645    }
646
647    // ── 23. Stats tracking — mixed forward and batch ──────────────────────────
648
649    #[test]
650    fn test_stats_mixed_forward_and_batch() {
651        let mut net = make_net(4, 8, 4, FeedForwardActivation::Linear);
652        net.forward(&[0.0; 4]); // +1 call, +1 token
653        net.forward_batch(&vec![vec![0.0; 4]; 3]); // +1 call, +3 tokens
654        assert_eq!(net.stats().total_forward_calls, 2);
655        assert_eq!(net.stats().total_tokens_processed, 4);
656    }
657
658    // ── 24. Zero input ────────────────────────────────────────────────────────
659
660    #[test]
661    fn test_zero_input_with_bias() {
662        // For zero input, output = activation( bias1 ) projected through layer2.
663        // With zero-initialised biases the hidden layer is all-zero, so after any
664        // activation the output should also be all-zero.
665        let mut net = make_net(4, 8, 4, FeedForwardActivation::ReLU);
666        let out = net.forward(&[0.0; 4]);
667        // All biases are zero, so hidden = 0 after ReLU = 0, output = all zeros.
668        for v in &out {
669            assert!((v).abs() < 1e-30, "expected 0, got {}", v);
670        }
671    }
672
673    // ── 25. Unit input ────────────────────────────────────────────────────────
674
675    #[test]
676    fn test_unit_input_finite() {
677        let mut net = make_net(8, 32, 8, FeedForwardActivation::SiLU);
678        let out = net.forward(&[1.0; 8]);
679        for v in &out {
680            assert!(v.is_finite());
681        }
682    }
683
684    // ── 26. Dimension mismatch — short input (graceful) ───────────────────────
685
686    #[test]
687    fn test_short_input_graceful() {
688        // Provide a shorter-than-expected input; should not panic.
689        let mut net = make_net(8, 16, 8, FeedForwardActivation::ReLU);
690        let out = net.forward(&[1.0; 4]); // only 4 of 8 expected inputs
691        assert_eq!(out.len(), 8, "output dim should still be 8");
692        for v in &out {
693            assert!(v.is_finite());
694        }
695    }
696
697    // ── 27. Dimension mismatch — zero-dim input ───────────────────────────────
698
699    #[test]
700    fn test_empty_input_graceful() {
701        let mut net = make_net(4, 8, 4, FeedForwardActivation::GELU);
702        let out = net.forward(&[]);
703        // With empty input all dot products are zero, so output == biases == 0
704        assert_eq!(out.len(), 4);
705        for v in &out {
706            assert!(v.is_finite());
707        }
708    }
709
710    // ── 28. No-bias layer ─────────────────────────────────────────────────────
711
712    #[test]
713    fn test_no_bias_zero_input_is_zero() {
714        let mut net = make_net_no_bias(4, 8, 4);
715        let out = net.forward(&[0.0; 4]);
716        for v in &out {
717            assert!(v.abs() < 1e-30);
718        }
719    }
720
721    // ── 29. Sigmoid basic properties ──────────────────────────────────────────
722
723    #[test]
724    fn test_sigmoid_properties() {
725        assert!((FeedForwardNetwork::sigmoid(0.0) - 0.5).abs() < 1e-12);
726        assert!(FeedForwardNetwork::sigmoid(100.0) > 0.999);
727        assert!(FeedForwardNetwork::sigmoid(-100.0) < 0.001);
728    }
729
730    // ── 30. FeedForwardActivation equality ───────────────────────────────────
731
732    #[test]
733    fn test_activation_enum_equality() {
734        assert_eq!(FeedForwardActivation::ReLU, FeedForwardActivation::ReLU);
735        assert_ne!(FeedForwardActivation::ReLU, FeedForwardActivation::GELU);
736        assert_ne!(FeedForwardActivation::SiLU, FeedForwardActivation::Linear);
737    }
738
739    // ── 31. Empty batch ───────────────────────────────────────────────────────
740
741    #[test]
742    fn test_empty_batch() {
743        let mut net = make_net(4, 8, 4, FeedForwardActivation::ReLU);
744        let out = net.forward_batch(&[]);
745        assert!(out.is_empty());
746    }
747
748    // ── 32. Weight matrix shape from init_layer ───────────────────────────────
749
750    #[test]
751    fn test_init_layer_shape() {
752        let mut rng = 1_u64;
753        let layer = FeedForwardNetwork::init_layer(6, 12, true, &mut rng);
754        assert_eq!(layer.weights.len(), 12, "out_dim rows");
755        for row in &layer.weights {
756            assert_eq!(row.len(), 6, "in_dim cols");
757        }
758        assert_eq!(layer.bias.len(), 12);
759    }
760}