optirs_learned/transformer/architecture/
encoder.rs

1use std::fmt::Debug;
2// Transformer encoder layers and components
3//
4// This module implements the encoder components of the transformer optimizer,
5// including the transformer layer, feed-forward network, and layer normalization.
6
7use scirs2_core::ndarray::{s, Array1, Array2};
8use scirs2_core::numeric::Float;
9use scirs2_core::random::CoreRandom as Random;
10
11use super::super::TransformerOptimizerConfig;
12use super::attention::MultiHeadAttention;
13use crate::error::{OptimError, Result};
14
15/// Activation functions for feed-forward networks
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
17pub enum ActivationFunction {
18    /// ReLU activation
19    ReLU,
20    /// GELU activation
21    #[default]
22    GELU,
23    /// Swish/SiLU activation
24    Swish,
25    /// GLU (Gated Linear Unit): splits the input and gates one half by
26    /// the sigmoid of the other, halving the width.
27    GLU,
28    /// GeGLU (GELU variant of GLU): splits the input and gates one half by
29    /// the GELU of the other, halving the width.
30    GeGLU,
31}
32
33impl ActivationFunction {
34    /// Whether this activation halves the feature width (gated variants do).
35    pub fn is_gated(self) -> bool {
36        matches!(self, ActivationFunction::GLU | ActivationFunction::GeGLU)
37    }
38
39    /// Width produced by this activation given `input_width`.
40    pub fn output_width(self, input_width: usize) -> usize {
41        if self.is_gated() {
42            input_width / 2
43        } else {
44            input_width
45        }
46    }
47}
48
49/// Single transformer encoder layer
50#[derive(Debug, Clone)]
51pub struct TransformerLayer<T: Float + Debug + Send + Sync + 'static> {
52    /// Multi-head self-attention
53    self_attention: MultiHeadAttention<T>,
54
55    /// Cross-attention (for multi-task learning)
56    cross_attention: Option<MultiHeadAttention<T>>,
57
58    /// Feed-forward network
59    feed_forward: FeedForwardNetwork<T>,
60
61    /// Layer normalization layers
62    ln1: LayerNorm<T>,
63    ln2: LayerNorm<T>,
64    ln3: Option<LayerNorm<T>>, // For cross-attention
65
66    /// Dropout layers
67    dropout1: DropoutLayer,
68    dropout2: DropoutLayer,
69    dropout3: Option<DropoutLayer>,
70
71    /// Use pre-layer normalization
72    pre_layer_norm: bool,
73}
74
75/// Feed-forward network
76#[derive(Debug, Clone)]
77pub struct FeedForwardNetwork<T: Float + Debug + Send + Sync + 'static> {
78    /// First linear layer weights
79    linear1: Array2<T>,
80
81    /// First linear layer bias
82    bias1: Array1<T>,
83
84    /// Second linear layer weights
85    linear2: Array2<T>,
86
87    /// Second linear layer bias
88    bias2: Array1<T>,
89
90    /// Activation function
91    activation: ActivationFunction,
92
93    /// Model dimension (input and output width)
94    modeldim: usize,
95
96    /// Hidden width produced by the first linear layer
97    ff_dim: usize,
98
99    /// Dropout layer
100    dropout: DropoutLayer,
101}
102
103/// Layer normalization
104#[derive(Debug, Clone)]
105pub struct LayerNorm<T: Float + Debug + Send + Sync + 'static> {
106    /// Scale parameters (gamma)
107    gamma: Array1<T>,
108
109    /// Shift parameters (beta)
110    beta: Array1<T>,
111
112    /// Epsilon for numerical stability
113    eps: T,
114
115    /// Dimension
116    dim: usize,
117}
118
119/// Dropout layer with a deterministic, seeded generator.
120#[derive(Debug, Clone)]
121pub struct DropoutLayer {
122    /// Dropout probability
123    prob: f64,
124
125    /// Training mode. Defaults to `false` so that inference is the safe default;
126    /// call [`DropoutLayer::set_training`] to enable dropout during training.
127    training: bool,
128
129    /// xorshift64* state (never zero)
130    rng_state: u64,
131}
132
133impl<
134        T: Float
135            + Debug
136            + Default
137            + Clone
138            + std::iter::Sum
139            + scirs2_core::ndarray::ScalarOperand
140            + Send
141            + Sync,
142    > TransformerLayer<T>
143{
144    /// Build a transformer encoder layer from the optimizer configuration.
145    pub fn new(config: &TransformerOptimizerConfig, _rng: &mut Random) -> Result<Self> {
146        let self_attention = MultiHeadAttention::new(config)?;
147        let cross_attention = if config.cross_attention {
148            Some(MultiHeadAttention::new(config)?)
149        } else {
150            None
151        };
152
153        let feed_forward = FeedForwardNetwork::new(config)?;
154
155        let ln1 = LayerNorm::new(config.modeldim, config.layer_norm_eps);
156        let ln2 = LayerNorm::new(config.modeldim, config.layer_norm_eps);
157        let ln3 = if config.cross_attention {
158            Some(LayerNorm::new(config.modeldim, config.layer_norm_eps))
159        } else {
160            None
161        };
162
163        let dropout1 = DropoutLayer::with_seed(config.attention_dropout, config.dropout_seed);
164        let dropout2 = DropoutLayer::with_seed(config.ff_dropout, config.dropout_seed ^ 0x9E37);
165        let dropout3 = if config.cross_attention {
166            Some(DropoutLayer::with_seed(
167                config.attention_dropout,
168                config.dropout_seed ^ 0x7F4A,
169            ))
170        } else {
171            None
172        };
173
174        Ok(Self {
175            self_attention,
176            cross_attention,
177            feed_forward,
178            ln1,
179            ln2,
180            ln3,
181            dropout1,
182            dropout2,
183            dropout3,
184            pre_layer_norm: config.pre_layer_norm,
185        })
186    }
187
188    /// Forward pass through the layer.
189    pub fn forward(&mut self, input: &Array2<T>) -> Result<Array2<T>> {
190        let mut x = input.clone();
191
192        // Self-attention with residual connection
193        let residual = x.clone();
194        if self.pre_layer_norm {
195            x = self.ln1.forward(&x)?;
196        }
197
198        x = self.self_attention.forward(&x, &x, &x)?;
199        x = self.dropout1.forward(&x);
200        x = x + &residual;
201
202        if !self.pre_layer_norm {
203            x = self.ln1.forward(&x)?;
204        }
205
206        // Cross-attention (if enabled)
207        if let Some(ref mut cross_attn) = self.cross_attention {
208            let residual = x.clone();
209            if self.pre_layer_norm {
210                if let Some(ref ln3) = self.ln3 {
211                    x = ln3.forward(&x)?;
212                }
213            }
214
215            x = cross_attn.forward(&x, &x, &x)?;
216            if let Some(ref mut dropout3) = self.dropout3 {
217                x = dropout3.forward(&x);
218            }
219            x = x + &residual;
220
221            if !self.pre_layer_norm {
222                if let Some(ref ln3) = self.ln3 {
223                    x = ln3.forward(&x)?;
224                }
225            }
226        }
227
228        // Feed-forward with residual connection
229        let residual = x.clone();
230        if self.pre_layer_norm {
231            x = self.ln2.forward(&x)?;
232        }
233
234        x = self.feed_forward.forward(&x)?;
235        x = self.dropout2.forward(&x);
236        x = x + &residual;
237
238        if !self.pre_layer_norm {
239            x = self.ln2.forward(&x)?;
240        }
241
242        Ok(x)
243    }
244
245    /// Get attention patterns for analysis
246    pub fn get_attention_patterns(&self) -> Option<&scirs2_core::ndarray::Array3<T>> {
247        self.self_attention.get_attention_patterns()
248    }
249
250    /// Switch every dropout in the layer between training and inference mode.
251    pub fn set_training(&mut self, training: bool) {
252        self.dropout1.set_training(training);
253        self.dropout2.set_training(training);
254        if let Some(ref mut dropout3) = self.dropout3 {
255            dropout3.set_training(training);
256        }
257        self.feed_forward.set_training(training);
258    }
259
260    /// Mutable access to the feed-forward network
261    pub fn feed_forward_mut(&mut self) -> &mut FeedForwardNetwork<T> {
262        &mut self.feed_forward
263    }
264
265    /// Access the self-attention block
266    pub fn self_attention(&self) -> &MultiHeadAttention<T> {
267        &self.self_attention
268    }
269}
270
271impl<T: Float + Debug + Default + Clone + Send + Sync + 'static> FeedForwardNetwork<T> {
272    /// Build a feed-forward network from the optimizer configuration.
273    pub fn new(config: &TransformerOptimizerConfig) -> Result<Self> {
274        Self::with_dimensions(
275            config.modeldim,
276            config.ff_dim,
277            config.activation,
278            config.ff_dropout,
279            config.dropout_seed ^ 0x1357,
280        )
281    }
282
283    /// Build a feed-forward network with explicit dimensions.
284    ///
285    /// The second linear layer is sized from the *activation output* width, so
286    /// gated activations (which halve the width) are handled correctly.
287    pub fn with_dimensions(
288        modeldim: usize,
289        ff_dim: usize,
290        activation: ActivationFunction,
291        dropout: f64,
292        dropout_seed: u64,
293    ) -> Result<Self> {
294        if modeldim == 0 || ff_dim == 0 {
295            return Err(OptimError::InvalidConfig(
296                "Feed-forward dimensions must be positive".to_string(),
297            ));
298        }
299        if activation.is_gated() && !ff_dim.is_multiple_of(2) {
300            return Err(OptimError::InvalidConfig(format!(
301                "Gated activation {activation:?} requires an even feed-forward dimension, got {ff_dim}"
302            )));
303        }
304
305        let hidden_out = activation.output_width(ff_dim);
306        let linear1 = Self::xavier(modeldim, ff_dim);
307        let linear2 = Self::xavier(hidden_out, modeldim);
308
309        Ok(Self {
310            linear1,
311            bias1: Array1::zeros(ff_dim),
312            linear2,
313            bias2: Array1::zeros(modeldim),
314            activation,
315            modeldim,
316            ff_dim,
317            dropout: DropoutLayer::with_seed(dropout, dropout_seed),
318        })
319    }
320
321    /// Xavier/Glorot uniform initialization with limit `sqrt(6 / (fan_in + fan_out))`.
322    fn xavier(fan_in: usize, fan_out: usize) -> Array2<T> {
323        let mut rng = scirs2_core::random::thread_rng();
324        let bound = (6.0 / (fan_in + fan_out).max(1) as f64).sqrt();
325        let mut weights = Array2::zeros((fan_in, fan_out));
326        for elem in weights.iter_mut() {
327            *elem = scirs2_core::numeric::NumCast::from((rng.random::<f64>() - 0.5) * 2.0 * bound)
328                .unwrap_or_else(|| T::zero());
329        }
330        weights
331    }
332
333    /// Forward pass through the feed-forward network.
334    pub fn forward(&mut self, input: &Array2<T>) -> Result<Array2<T>> {
335        let x1 = Self::linear_transform(input, &self.linear1, &self.bias1)?;
336        let x2 = self.apply_activation(&x1);
337        let x3 = self.dropout.forward(&x2);
338        Self::linear_transform(&x3, &self.linear2, &self.bias2)
339    }
340
341    fn linear_transform(
342        input: &Array2<T>,
343        weights: &Array2<T>,
344        bias: &Array1<T>,
345    ) -> Result<Array2<T>> {
346        let (_seq_len, input_dim) = input.dim();
347        let (weight_in, weight_out) = weights.dim();
348
349        if input_dim != weight_in {
350            return Err(OptimError::InvalidConfig(format!(
351                "Input dimension {input_dim} doesn't match weight matrix input {weight_in}"
352            )));
353        }
354        if bias.len() != weight_out {
355            return Err(OptimError::InvalidConfig(format!(
356                "Bias dimension {} doesn't match output dimension {weight_out}",
357                bias.len()
358            )));
359        }
360
361        Ok(input.dot(weights) + bias)
362    }
363
364    /// Numerically stable GELU (tanh approximation).
365    fn gelu(x: T) -> T {
366        let sqrt_2_pi: T =
367            scirs2_core::numeric::NumCast::from(0.797_884_560_802_865).unwrap_or_else(|| T::one());
368        let coeff: T = scirs2_core::numeric::NumCast::from(0.044715).unwrap_or_else(|| T::zero());
369        let half: T = scirs2_core::numeric::NumCast::from(0.5).unwrap_or_else(|| T::one());
370        let inner = sqrt_2_pi * (x + coeff * x * x * x);
371        half * x * (T::one() + inner.tanh())
372    }
373
374    /// Numerically stable logistic sigmoid.
375    fn sigmoid(x: T) -> T {
376        if x >= T::zero() {
377            T::one() / (T::one() + (-x).exp())
378        } else {
379            let e = x.exp();
380            e / (T::one() + e)
381        }
382    }
383
384    fn apply_activation(&self, input: &Array2<T>) -> Array2<T> {
385        match self.activation {
386            ActivationFunction::ReLU => input.map(|&x| if x > T::zero() { x } else { T::zero() }),
387            ActivationFunction::GELU => input.map(|&x| Self::gelu(x)),
388            // Swish / SiLU: x * sigmoid(x). The previous `x * exp(x) / (1 +
389            // exp(x))` form overflowed to NaN for x greater than about 710.
390            ActivationFunction::Swish => input.map(|&x| x * Self::sigmoid(x)),
391            ActivationFunction::GLU => self.gated(input, Self::sigmoid),
392            ActivationFunction::GeGLU => self.gated(input, Self::gelu),
393        }
394    }
395
396    /// Gated linear unit: split the last axis in half and gate `a` by `g(b)`.
397    fn gated(&self, input: &Array2<T>, gate: fn(T) -> T) -> Array2<T> {
398        let (rows, cols) = input.dim();
399        let half = cols / 2;
400        let mut output = Array2::zeros((rows, half));
401        for i in 0..rows {
402            for j in 0..half {
403                output[[i, j]] = input[[i, j]] * gate(input[[i, half + j]]);
404            }
405        }
406        output
407    }
408
409    /// Set the activation function, resizing the second linear layer when the
410    /// activation changes the hidden width (gated vs. elementwise).
411    pub fn set_activation(&mut self, activation: ActivationFunction) -> Result<()> {
412        if activation == self.activation {
413            return Ok(());
414        }
415        if activation.is_gated() && !self.ff_dim.is_multiple_of(2) {
416            return Err(OptimError::InvalidConfig(format!(
417                "Gated activation {activation:?} requires an even feed-forward dimension, got {}",
418                self.ff_dim
419            )));
420        }
421
422        let new_width = activation.output_width(self.ff_dim);
423        if new_width != self.linear2.nrows() {
424            self.linear2 = Self::xavier(new_width, self.modeldim);
425        }
426        self.activation = activation;
427        Ok(())
428    }
429
430    /// Get the active activation function
431    pub fn activation(&self) -> ActivationFunction {
432        self.activation
433    }
434
435    /// Model (input/output) dimension
436    pub fn model_dim(&self) -> usize {
437        self.modeldim
438    }
439
440    /// Hidden dimension of the first linear layer
441    pub fn hidden_dim(&self) -> usize {
442        self.ff_dim
443    }
444
445    /// Enable or disable dropout
446    pub fn set_training(&mut self, training: bool) {
447        self.dropout.set_training(training);
448    }
449}
450
451impl<T: Float + Debug + Default + Clone + std::iter::Sum + Send + Sync> LayerNorm<T> {
452    /// Create a layer normalization block with the given epsilon.
453    pub fn new(dim: usize, eps: f64) -> Self {
454        Self {
455            gamma: Array1::ones(dim),
456            beta: Array1::zeros(dim),
457            eps: scirs2_core::numeric::NumCast::from(eps).unwrap_or_else(|| T::zero()),
458            dim,
459        }
460    }
461
462    /// Create a layer normalization block with the default epsilon of 1e-6.
463    pub fn with_default_eps(dim: usize) -> Self {
464        Self::new(dim, 1e-6)
465    }
466
467    /// Normalize each row: `(x - mean) / sqrt(var + eps) * gamma + beta`.
468    pub fn forward(&self, input: &Array2<T>) -> Result<Array2<T>> {
469        let (seq_len, input_dim) = input.dim();
470
471        if input_dim != self.dim {
472            return Err(OptimError::InvalidConfig(format!(
473                "Input dimension {} doesn't match layer norm dimension {}",
474                input_dim, self.dim
475            )));
476        }
477        if input_dim == 0 {
478            return Ok(Array2::zeros((seq_len, 0)));
479        }
480
481        let mut output = Array2::zeros((seq_len, input_dim));
482        let width: T =
483            scirs2_core::numeric::NumCast::from(input_dim as f64).unwrap_or_else(|| T::one());
484
485        for i in 0..seq_len {
486            let row = input.slice(s![i, ..]);
487
488            let mean = row.iter().cloned().sum::<T>() / width;
489            let variance = row
490                .iter()
491                .map(|&x| {
492                    let diff = x - mean;
493                    diff * diff
494                })
495                .sum::<T>()
496                / width;
497
498            let std = (variance + self.eps).sqrt();
499
500            for j in 0..input_dim {
501                let normalized = (input[[i, j]] - mean) / std;
502                output[[i, j]] = self.gamma[j] * normalized + self.beta[j];
503            }
504        }
505
506        Ok(output)
507    }
508
509    /// Get layer normalization parameters
510    pub fn parameters(&self) -> (&Array1<T>, &Array1<T>) {
511        (&self.gamma, &self.beta)
512    }
513
514    /// Normalized dimension
515    pub fn dim(&self) -> usize {
516        self.dim
517    }
518
519    /// Set layer normalization parameters
520    pub fn set_parameters(&mut self, gamma: Array1<T>, beta: Array1<T>) -> Result<()> {
521        if gamma.len() != self.dim || beta.len() != self.dim {
522            return Err(OptimError::InvalidConfig(format!(
523                "Parameter dimensions ({}, {}) don't match layer norm dimension {}",
524                gamma.len(),
525                beta.len(),
526                self.dim
527            )));
528        }
529        self.gamma = gamma;
530        self.beta = beta;
531        Ok(())
532    }
533}
534
535impl DropoutLayer {
536    /// Create a dropout layer with a default seed. Starts in inference mode.
537    pub fn new(prob: f64) -> Self {
538        Self::with_seed(prob, 0x243F_6A88_85A3_08D3)
539    }
540
541    /// Create a dropout layer with an explicit seed. Starts in inference mode.
542    pub fn with_seed(prob: f64, seed: u64) -> Self {
543        Self {
544            prob: prob.clamp(0.0, 1.0),
545            training: false,
546            // xorshift64* requires a non-zero state.
547            rng_state: seed | 1,
548        }
549    }
550
551    fn next_f64(&mut self) -> f64 {
552        let mut x = self.rng_state;
553        x ^= x >> 12;
554        x ^= x << 25;
555        x ^= x >> 27;
556        self.rng_state = x;
557        let value = x.wrapping_mul(0x2545_F491_4F6C_DD1D);
558        (value >> 11) as f64 / (1u64 << 53) as f64
559    }
560
561    /// Inverted dropout: zero each entry with probability `prob` and rescale the
562    /// survivors by `1 / (1 - prob)` so the expected activation is unchanged.
563    /// In inference mode the input is returned untouched.
564    pub fn forward<T: Float + Clone>(&mut self, input: &Array2<T>) -> Array2<T> {
565        if !self.training || self.prob <= 0.0 {
566            return input.clone();
567        }
568        if self.prob >= 1.0 {
569            return Array2::zeros(input.dim());
570        }
571
572        let keep_prob = 1.0 - self.prob;
573        let scale: T = scirs2_core::numeric::NumCast::from(1.0 / keep_prob).unwrap_or_else(T::one);
574
575        let mut output = input.clone();
576        for elem in output.iter_mut() {
577            if self.next_f64() < self.prob {
578                *elem = T::zero();
579            } else {
580                *elem = *elem * scale;
581            }
582        }
583        output
584    }
585
586    /// Set training mode
587    pub fn set_training(&mut self, training: bool) {
588        self.training = training;
589    }
590
591    /// Whether the layer is in training mode
592    pub fn is_training(&self) -> bool {
593        self.training
594    }
595
596    /// Get dropout probability
597    pub fn prob(&self) -> f64 {
598        self.prob
599    }
600}
601
602#[cfg(test)]
603mod tests {
604    use super::*;
605
606    #[test]
607    fn swish_is_stable_for_large_inputs() {
608        let input =
609            Array2::<f64>::from_shape_vec((1, 3), vec![-1000.0, 0.0, 1000.0]).expect("valid shape");
610        let mut ffn =
611            FeedForwardNetwork::<f64>::with_dimensions(3, 4, ActivationFunction::Swish, 0.0, 1)
612                .expect("ffn creation");
613        let activated = ffn.apply_activation(&input);
614        assert!(activated.iter().all(|v| v.is_finite()), "{activated:?}");
615        assert!((activated[[0, 2]] - 1000.0).abs() < 1e-6);
616        assert!(activated[[0, 0]].abs() < 1e-6);
617        ffn.set_training(true);
618    }
619
620    #[test]
621    fn gated_activations_halve_the_width() {
622        let ffn = FeedForwardNetwork::<f64>::with_dimensions(4, 8, ActivationFunction::GLU, 0.0, 1)
623            .expect("ffn creation");
624        let hidden = Array2::<f64>::from_shape_fn((2, 8), |(_, j)| j as f64 - 3.5);
625        let activated = ffn.apply_activation(&hidden);
626        assert_eq!(activated.shape(), &[2, 4]);
627        // GLU: a * sigmoid(b)
628        let expected = hidden[[0, 0]] * (1.0 / (1.0 + (-hidden[[0, 4]]).exp()));
629        assert!((activated[[0, 0]] - expected).abs() < 1e-12);
630    }
631
632    #[test]
633    fn gated_feed_forward_round_trips_to_model_dim() {
634        let mut ffn =
635            FeedForwardNetwork::<f64>::with_dimensions(4, 8, ActivationFunction::GeGLU, 0.0, 1)
636                .expect("ffn creation");
637        let input = Array2::<f64>::from_shape_fn((3, 4), |(i, j)| (i + j) as f64 * 0.1);
638        let output = ffn.forward(&input).expect("forward");
639        assert_eq!(output.shape(), &[3, 4]);
640    }
641
642    #[test]
643    fn switching_to_a_gated_activation_resizes_the_output_layer() {
644        let mut ffn =
645            FeedForwardNetwork::<f64>::with_dimensions(4, 8, ActivationFunction::GELU, 0.0, 1)
646                .expect("ffn creation");
647        ffn.set_activation(ActivationFunction::GLU)
648            .expect("activation switch");
649        let input = Array2::<f64>::ones((2, 4));
650        let output = ffn.forward(&input).expect("forward");
651        assert_eq!(output.shape(), &[2, 4]);
652    }
653
654    #[test]
655    fn odd_hidden_width_rejects_gated_activation() {
656        let result =
657            FeedForwardNetwork::<f64>::with_dimensions(4, 7, ActivationFunction::GLU, 0.0, 1);
658        assert!(result.is_err());
659    }
660
661    #[test]
662    fn dropout_is_identity_at_inference_and_active_in_training() {
663        let mut dropout = DropoutLayer::with_seed(0.5, 42);
664        let input = Array2::<f64>::ones((8, 8));
665
666        // Default is inference mode.
667        assert!(!dropout.is_training());
668        assert_eq!(dropout.forward(&input), input);
669
670        dropout.set_training(true);
671        let output = dropout.forward(&input);
672        let zeros = output.iter().filter(|&&v| v == 0.0).count();
673        assert!(zeros > 0, "training-mode dropout dropped nothing");
674        assert!(
675            zeros < input.len(),
676            "training-mode dropout dropped everything"
677        );
678        // Inverted dropout rescales survivors by 1 / (1 - p) = 2.
679        assert!(output.iter().all(|&v| v == 0.0 || (v - 2.0).abs() < 1e-12));
680    }
681
682    #[test]
683    fn dropout_is_reproducible_for_a_given_seed() {
684        let input = Array2::<f64>::ones((4, 4));
685        let run = || {
686            let mut dropout = DropoutLayer::with_seed(0.3, 7);
687            dropout.set_training(true);
688            dropout.forward(&input)
689        };
690        assert_eq!(run(), run());
691    }
692
693    #[test]
694    fn layer_norm_rejects_mismatched_width() {
695        let norm = LayerNorm::<f64>::with_default_eps(8);
696        let input = Array2::<f64>::zeros((2, 4));
697        assert!(norm.forward(&input).is_err());
698    }
699
700    #[test]
701    fn layer_norm_standardizes_rows() {
702        let norm = LayerNorm::<f64>::with_default_eps(4);
703        let input =
704            Array2::<f64>::from_shape_vec((1, 4), vec![1.0, 2.0, 3.0, 4.0]).expect("valid shape");
705        let output = norm.forward(&input).expect("forward");
706        let mean: f64 = output.row(0).iter().sum::<f64>() / 4.0;
707        assert!(mean.abs() < 1e-9, "mean {mean}");
708    }
709}
optirs_learned/transformer/architecture/encoder.rs

optirs_learned/transformer/architecture/
encoder.rs
