## Coverage Report

#### Created: 2026-01-25 15:05

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /home/noah/src/realizar/src/layers/position.rs

    Line
    Count
    Source
    1
    //! Position embeddings for transformer models
    2
    //!
    3
    //! Extracted from layers/mod.rs (PMAT-802) to reduce module size.
    4
    //! Contains:
    5
    //! - RoPE: Rotary Position Embeddings (RoFormer, LLaMA, PaLM)
    6
    //! - RopeScalingType: Context length extension methods (NTK, YaRN, Linear)
    7
    //! - ScaledRoPE: RoPE with scaling for extended context
    8
    //! - ALiBi: Attention with Linear Biases
    9
    10
    use crate::{
    11
        error::{RealizarError, Result},
    12
        tensor::Tensor,
    13
    };
    14
    15
    /// Rotary Position Embeddings (`RoPE`)
    16
    ///
    17
    /// Applies position-dependent rotations to query and key vectors.
    18
    /// Used in `LLaMA`, `PaLM`, and other modern transformers for relative
    19
    /// position encoding.
    20
    ///
    21
    /// The rotation is applied pairwise to dimensions, encoding position
    22
    /// information directly into the embeddings.
    23
    ///
    24
    /// # Formula
    25
    ///
    26
    /// For each pair of dimensions (2i, 2i+1):
    27
    /// ```text
    28
    /// x'_{2i} = x_{2i} * cos(θ_i * pos) - x_{2i+1} * sin(θ_i * pos)
    29
    /// x'_{2i+1} = x_{2i} * sin(θ_i * pos) + x_{2i+1} * cos(θ_i * pos)
    30
    /// ```
    31
    ///
    32
    /// Where `θ_i` = base^(-2i/dim)
    33
    ///
    34
    /// # References
    35
    ///
    36
    /// `RoFormer`: Enhanced Transformer with Rotary Position Embedding - Su et al., 2021
    37
    #[derive(Debug, Clone)]
    38
    pub struct RoPE {
    39
        /// Embedding dimension (must be even)
    40
        dim: usize,
    41
        /// Base for computing frequencies (default: 10000)
    42
        base: f32,
    43
        /// Precomputed inverse frequencies for each dimension pair
    44
        inv_freq: Vec<f32>,
    45
    }
    46
    47
    impl RoPE {
    48
        /// Create a new `RoPE` layer
    49
        ///
    50
        /// # Arguments
    51
        ///
    52
        /// * `dim` - Embedding dimension (must be even)
    53
        /// * `base` - Base for computing frequencies (typically 10000)
    54
        ///
    55
        /// # Errors
    56
        ///
    57
        /// Returns error if `dim` is zero or odd
    58
    18
        pub fn new(dim: usize, base: f32) -> Result<Self> {
    59
    18
            if dim == 0 {
    60
    2
                return Err(RealizarError::InvalidShape {
    61
    2
                    reason: "dim must be > 0".to_string(),
    62
    2
                });
    63
    16
            }
    64
    16
            if !dim.is_multiple_of(2) {
    65
    2
                return Err(RealizarError::InvalidShape {
    66
    2
                    reason: "dim must be even for RoPE".to_string(),
    67
    2
                });
    68
    14
            }
    69
    70
            // Compute inverse frequencies: base^(-2i/dim) for i in 0..dim/2
    71
    14
            let half_dim = dim / 2;
    72
    14
            let mut inv_freq = Vec::with_capacity(half_dim);
    73
    74
            #[allow(clippy::cast_precision_loss)]
    75
    270
            for i in 0..half_dim14 {
    76
    270
                let exponent = -2.0 * (i as f32) / (dim as f32);
    77
    270
                inv_freq.push(base.powf(exponent));
    78
    270
            }
    79
    80
    14
            Ok(Self {
    81
    14
                dim,
    82
    14
                base,
    83
    14
                inv_freq,
    84
    14
            })
    85
    18
        }
    86
    87
        /// Create `RoPE` with default base (10000)
    88
        ///
    89
        /// # Errors
    90
        ///
    91
        /// Returns error if `dim` is zero or odd
    92
    7
        pub fn with_default_base(dim: usize) -> Result<Self> {
    93
    7
            Self::new(dim, 10000.0)
    94
    7
        }
    95
    96
        /// Apply rotary embeddings to input at given position
    97
        ///
    98
        /// # Arguments
    99
        ///
    100
        /// * `input` - Input tensor with last dimension equal to `dim`
    101
        /// * `position` - Position index for computing rotation angles
    102
        ///
    103
        /// # Returns
    104
        ///
    105
        /// Tensor with same shape as input, with rotary embeddings applied
    106
        ///
    107
        /// # Errors
    108
        ///
    109
        /// Returns error if input's last dimension doesn't match `dim`
    110
    10
        pub fn forward(&self, input: &Tensor<f32>, position: usize) -> Result<Tensor<f32>> {
    111
    10
            let shape = input.shape();
    112
    113
    10
            if shape.is_empty() {
    114
    0
                return Err(RealizarError::InvalidShape {
    115
    0
                    reason: "Input tensor must have at least 1 dimension".to_string(),
    116
    0
                });
    117
    10
            }
    118
    119
    10
            let last_dim = shape[shape.len() - 1];
    120
    10
            if last_dim != self.dim {
    121
    1
                return Err(RealizarError::InvalidShape {
    122
    1
                    reason: format!("Expected last dimension {}, got {}", self.dim, last_dim),
    123
    1
                });
    124
    9
            }
    125
    126
    9
            let data = input.data();
    127
    9
            let num_vectors = data.len() / self.dim;
    128
    9
            let mut output = Vec::with_capacity(data.len());
    129
    130
            // Compute sin/cos for this position
    131
    9
            let half_dim = self.dim / 2;
    132
    9
            let mut cos_vals = Vec::with_capacity(half_dim);
    133
    9
            let mut sin_vals = Vec::with_capacity(half_dim);
    134
    135
            #[allow(clippy::cast_precision_loss)]
    136
    87
            for inv_f78 in &self.inv_freq {
    137
    78
                let angle = inv_f * (position as f32);
    138
    78
                cos_vals.push(angle.cos());
    139
    78
                sin_vals.push(angle.sin());
    140
    78
            }
    141
    142
            // Apply rotation to each vector
    143
    12
            for vec_idx in 0..num_vectors9 {
    144
    12
                let offset = vec_idx * self.dim;
    145
    146
    84
                for i in 0..half_dim12 {
    147
    84
                    let x0 = data[offset + 2 * i];
    148
    84
                    let x1 = data[offset + 2 * i + 1];
    149
    84
                    let cos_val = cos_vals[i];
    150
    84
                    let sin_val = sin_vals[i];
    151
    84
    152
    84
                    // Apply 2D rotation
    153
    84
                    let y0 = x0 * cos_val - x1 * sin_val;
    154
    84
                    let y1 = x0 * sin_val + x1 * cos_val;
    155
    84
    156
    84
                    output.push(y0);
    157
    84
                    output.push(y1);
    158
    84
                }
    159
            }
    160
    161
    9
            Tensor::from_vec(shape.to_vec(), output)
    162
    10
        }
    163
    164
        /// Get embedding dimension
    165
        #[must_use]
    166
    5
        pub fn dim(&self) -> usize {
    167
    5
            self.dim
    168
    5
        }
    169
    170
        /// Get base frequency
    171
        #[must_use]
    172
    3
        pub fn base(&self) -> f32 {
    173
    3
            self.base
    174
    3
        }
    175
    176
        /// Get inverse frequencies
    177
        #[must_use]
    178
    2
        pub fn inv_freq(&self) -> &[f32] {
    179
    2
            &self.inv_freq
    180
    2
        }
    181
    }
    182
    183
    // ============================================================================
    184
    // RoPE Scaling Methods (NTK, YaRN, Linear, Dynamic NTK)
    185
    // ============================================================================
    186
    //
    187
    // These methods extend RoPE to handle longer context lengths than trained.
    188
    // References:
    189
    // - NTK-aware: https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/
    190
    // - YaRN: https://arxiv.org/abs/2309.00071
    191
    // - Code Llama linear scaling: https://arxiv.org/abs/2308.12950
    192
    // ============================================================================
    193
    194
    /// RoPE scaling type for context length extension
    195
    #[derive(Debug, Clone, Copy, PartialEq, Default)]
    196
    pub enum RopeScalingType {
    197
        /// No scaling (original RoPE)
    198
        #[default]
    199
        None,
    200
        /// Linear interpolation (Code Llama style)
    201
        /// scale = trained_length / target_length
    202
        Linear {
    203
            /// Scale factor (typically trained_length / target_length)
    204
            scale: f32,
    205
        },
    206
        /// NTK-aware scaling
    207
        /// Modifies base frequency: base' = base * scale^(dim / (dim - 2))
    208
        Ntk {
    209
            /// Scale factor for context extension
    210
            scale: f32,
    211
        },
    212
        /// Dynamic NTK-aware scaling
    213
        /// Adjusts scale dynamically based on current sequence length
    214
        DynamicNtk {
    215
            /// Original training context length
    216
            original_max_len: usize,
    217
            /// Target extended context length
    218
            target_max_len: usize,
    219
        },
    220
        /// YaRN (Yet another RoPE extensioN)
    221
        /// Combines NTK interpolation with attention scaling
    222
        Yarn {
    223
            /// Original training context length
    224
            original_max_len: usize,
    225
            /// Target extended context length
    226
            target_max_len: usize,
    227
            /// Attention scaling factor (typically sqrt(scale))
    228
            attn_factor: f32,
    229
            /// Beta for interpolation ramp (default: 32)
    230
            beta_fast: f32,
    231
            /// Beta for extrapolation (default: 1)
    232
            beta_slow: f32,
    233
        },
    234
    }
    235
    236
    /// Scaled Rotary Position Embeddings
    237
    ///
    238
    /// Extends `RoPE` with various scaling methods for context length extension.
    239
    /// Supports NTK-aware, Linear, Dynamic NTK, and YaRN scaling.
    240
    ///
    241
    /// # Scaling Methods
    242
    ///
    243
    /// ## Linear Scaling (Code Llama)
    244
    /// Simply scales down the position: pos' = pos / scale
    245
    ///
    246
    /// ## NTK-aware Scaling
    247
    /// Modifies the base frequency to reduce high-frequency component decay:
    248
    /// base' = base * scale^(dim / (dim - 2))
    249
    ///
    250
    /// ## Dynamic NTK
    251
    /// Dynamically adjusts NTK scale based on current sequence length
    252
    ///
    253
    /// ## YaRN (Yet another RoPE extensioN)
    254
    /// Combines NTK with attention factor and interpolation ramp
    255
    ///
    256
    /// # References
    257
    ///
    258
    /// - "Code Llama: Open Foundation Models for Code" - Rozière et al., 2023
    259
    /// - "YaRN: Efficient Context Window Extension of Large Language Models" - Peng et al., 2023
    260
    #[derive(Debug, Clone)]
    261
    pub struct ScaledRoPE {
    262
        /// Base RoPE parameters
    263
        dim: usize,
    264
        /// Original base frequency
    265
        original_base: f32,
    266
        /// Scaled base frequency (after NTK adjustment)
    267
        scaled_base: f32,
    268
        /// Scaling configuration
    269
        scaling: RopeScalingType,
    270
        /// Precomputed inverse frequencies (with scaling applied)
    271
        inv_freq: Vec<f32>,
    272
        /// Attention scaling factor (for YaRN)
    273
        mscale: f32,
    274
    }
    275
    276
    impl ScaledRoPE {
    277
        /// Create a new scaled `RoPE` layer
    278
        ///
    279
        /// # Arguments
    280
        ///
    281
        /// * `dim` - Embedding dimension (must be even)
    282
        /// * `base` - Base frequency (typically 10000)
    283
        /// * `scaling` - Scaling method to use
    284
        ///
    285
        /// # Errors
    286
        ///
    287
        /// Returns error if `dim` is zero or odd
    288
    19
        pub fn new(dim: usize, base: f32, scaling: RopeScalingType) -> Result<Self> {
    289
    19
            if dim == 0 {
    290
    1
                return Err(RealizarError::InvalidShape {
    291
    1
                    reason: "dim must be > 0".to_string(),
    292
    1
                });
    293
    18
            }
    294
    18
            if !dim.is_multiple_of(2) {
    295
    1
                return Err(RealizarError::InvalidShape {
    296
    1
                    reason: "dim must be even for RoPE".to_string(),
    297
    1
                });
    298
    17
            }
    299
    300
    17
            let (scaled_base, mscale, inv_freq) = Self::compute_frequencies(dim, base, &scaling);
    301
    302
    17
            Ok(Self {
    303
    17
                dim,
    304
    17
                original_base: base,
    305
    17
                scaled_base,
    306
    17
                scaling,
    307
    17
                inv_freq,
    308
    17
                mscale,
    309
    17
            })
    310
    19
        }
    311
    312
        /// Create scaled `RoPE` with default base (10000)
    313
        ///
    314
        /// # Errors
    315
        ///
    316
        /// Returns error if `dim` is zero or odd
    317
    1
        pub fn with_default_base(dim: usize, scaling: RopeScalingType) -> Result<Self> {
    318
    1
            Self::new(dim, 10000.0, scaling)
    319
    1
        }
    320
    321
        /// Compute inverse frequencies with scaling applied
    322
    17
        fn compute_frequencies(
    323
    17
            dim: usize,
    324
    17
            base: f32,
    325
    17
            scaling: &RopeScalingType,
    326
    17
        ) -> (f32, f32, Vec<f32>) {
    327
    17
            let half_dim = dim / 2;
    328
    329
            // Compute scaled base and mscale based on scaling type
    330
            #[allow(clippy::cast_precision_loss)]
    331
    17
            let (scaled_base, mscale) = match scaling {
    332
    10
                RopeScalingType::None | RopeScalingType::Linear { .. } => (base, 1.0),
    333
    3
                RopeScalingType::Ntk { scale } => {
    334
                    // NTK formula: base' = base * scale^(dim / (dim - 2))
    335
    3
                    let dim_f = dim as f32;
    336
    3
                    let exponent = dim_f / (dim_f - 2.0);
    337
    3
                    let ntk_base = base * scale.powf(exponent);
    338
    3
                    (ntk_base, 1.0)
    339
                },
    340
                RopeScalingType::DynamicNtk {
    341
    1
                    original_max_len,
    342
    1
                    target_max_len,
    343
                } => {
    344
                    // Dynamic NTK uses scale = target / original
    345
    1
                    let scale = (*target_max_len as f32) / (*original_max_len as f32);
    346
    1
                    let dim_f = dim as f32;
    347
    1
                    let exponent = dim_f / (dim_f - 2.0);
    348
    1
                    let ntk_base = base * scale.powf(exponent);
    349
    1
                    (ntk_base, 1.0)
    350
                },
    351
                RopeScalingType::Yarn {
    352
    3
                    original_max_len,
    353
    3
                    target_max_len,
    354
    3
                    attn_factor,
    355
    3
                    beta_fast,
    356
    3
                    beta_slow,
    357
                } => {
    358
                    // YaRN combines NTK with attention scaling
    359
    3
                    let scale = (*target_max_len as f32) / (*original_max_len as f32);
    360
    3
                    let dim_f = dim as f32;
    361
    362
                    // Compute NTK-style base modification
    363
                    // YaRN uses a smoother interpolation based on frequency
    364
    3
                    let exponent = dim_f / (dim_f - 2.0);
    365
    3
                    let ntk_base = base * scale.powf(exponent);
    366
    367
                    // Compute mscale (attention factor)
    368
                    // Default: sqrt(1 + ln(scale) / ln(original_max_len))
    369
    3
                    let mscale = if *attn_factor > 0.0 {
    370
    2
                        *attn_factor
    371
                    } else {
    372
    1
                        let log_scale = scale.ln();
    373
    1
                        let log_orig = (*original_max_len as f32).ln();
    374
    1
                        (1.0 + log_scale / log_orig).sqrt()
    375
                    };
    376
    377
                    // The beta parameters affect the interpolation ramp
    378
                    // but are applied per-frequency in the forward pass
    379
    3
                    let _ = (beta_fast, beta_slow); // Used in forward
    380
    381
    3
                    (ntk_base, mscale)
    382
                },
    383
            };
    384
    385
            // Compute inverse frequencies with scaled base
    386
    17
            let mut inv_freq = Vec::with_capacity(half_dim);
    387
    388
            #[allow(clippy::cast_precision_loss)]
    389
    426
            for i in 0..half_dim17 {
    390
    426
                let exponent = -2.0 * (i as f32) / (dim as f32);
    391
    426
                inv_freq.push(scaled_base.powf(exponent));
    392
    426
            }
    393
    394
    17
            (scaled_base, mscale, inv_freq)
    395
    17
        }
    396
    397
        /// Apply scaled rotary embeddings to input at given position
    398
        ///
    399
        /// # Arguments
    400
        ///
    401
        /// * `input` - Input tensor with last dimension equal to `dim`
    402
        /// * `position` - Position index for computing rotation angles
    403
        ///
    404
        /// # Returns
    405
        ///
    406
        /// Tensor with same shape as input, with scaled rotary embeddings applied
    407
        ///
    408
        /// # Errors
    409
        ///
    410
        /// Returns error if input's last dimension doesn't match `dim`
    411
    5
        pub fn forward(&self, input: &Tensor<f32>, position: usize) -> Result<Tensor<f32>> {
    412
    5
            let shape = input.shape();
    413
    414
    5
            if shape.is_empty() {
    415
    0
                return Err(RealizarError::InvalidShape {
    416
    0
                    reason: "Input tensor must have at least 1 dimension".to_string(),
    417
    0
                });
    418
    5
            }
    419
    420
    5
            let last_dim = shape[shape.len() - 1];
    421
    5
            if last_dim != self.dim {
    422
    1
                return Err(RealizarError::InvalidShape {
    423
    1
                    reason: format!("Expected last dimension {}, got {}", self.dim, last_dim),
    424
    1
                });
    425
    4
            }
    426
    427
    4
            let data = input.data();
    428
    4
            let num_vectors = data.len() / self.dim;
    429
    4
            let mut output = Vec::with_capacity(data.len());
    430
    431
            // Compute effective position based on scaling type
    432
            #[allow(clippy::cast_precision_loss)]
    433
    4
            let effective_pos = match &self.scaling {
    434
                RopeScalingType::None
    435
                | RopeScalingType::Ntk { .. }
    436
                | RopeScalingType::DynamicNtk { .. }
    437
    3
                | RopeScalingType::Yarn { .. } => position as f32,
    438
    1
                RopeScalingType::Linear { scale } => (position as f32) / scale,
    439
            };
    440
    441
            // Compute sin/cos for this position
    442
    4
            let half_dim = self.dim / 2;
    443
    4
            let mut cos_vals = Vec::with_capacity(half_dim);
    444
    4
            let mut sin_vals = Vec::with_capacity(half_dim);
    445
    446
            // Apply YaRN interpolation ramp if applicable
    447
            #[allow(clippy::cast_precision_loss)]
    448
    8
            for (i, inv_f) in self.inv_freq.iter()4.enumerate4() {
    449
    8
                let angle = inv_f * effective_pos;
    450
    451
                // For YaRN, apply interpolation ramp based on frequency index
    452
    8
                let (cos_val, sin_val) = if let RopeScalingType::Yarn {
    453
    2
                    original_max_len,
    454
    2
                    target_max_len,
    455
    2
                    beta_fast,
    456
    2
                    beta_slow,
    457
                    ..
    458
    8
                } = &self.scaling
    459
                {
    460
                    // Compute wavelength for this frequency
    461
    2
                    let freq = 1.0 / inv_f;
    462
    2
                    let wavelength = 2.0 * std::f32::consts::PI * freq;
    463
    464
                    // Compute interpolation factor
    465
    2
                    let low_freq_wavelen = (*original_max_len as f32) / *beta_slow;
    466
    2
                    let high_freq_wavelen = (*original_max_len as f32) / *beta_fast;
    467
    468
    2
                    let ramp = if wavelength < high_freq_wavelen {
    469
    1
                        0.0 // Full extrapolation (use NTK-scaled frequency)
    470
    1
                    } else if wavelength > low_freq_wavelen {
    471
    1
                        1.0 // Full interpolation (use linear-scaled position)
    472
                    } else {
    473
                        // Linear ramp between extrapolation and interpolation
    474
    0
                        (wavelength - high_freq_wavelen) / (low_freq_wavelen - high_freq_wavelen)
    475
                    };
    476
    477
                    // Interpolate between NTK angle and linear angle
    478
    2
                    let scale = (*target_max_len as f32) / (*original_max_len as f32);
    479
    2
                    let linear_pos = effective_pos / scale;
    480
    481
                    // Original frequency from unscaled base
    482
    2
                    let orig_inv_f = self
    483
    2
                        .original_base
    484
    2
                        .powf(-2.0 * (i as f32) / (self.dim as f32));
    485
    2
                    let linear_angle = orig_inv_f * linear_pos;
    486
    487
                    // Interpolate angles
    488
    2
                    let final_angle = angle * (1.0 - ramp) + linear_angle * ramp;
    489
    490
    2
                    (final_angle.cos(), final_angle.sin())
    491
                } else {
    492
    6
                    (angle.cos(), angle.sin())
    493
                };
    494
    495
    8
                cos_vals.push(cos_val);
    496
    8
                sin_vals.push(sin_val);
    497
            }
    498
    499
            // Apply rotation to each vector (with mscale for YaRN)
    500
    4
            for vec_idx in 0..num_vectors {
    501
    4
                let offset = vec_idx * self.dim;
    502
    503
    8
                for i in 0..half_dim4 {
    504
    8
                    let x0 = data[offset + 2 * i];
    505
    8
                    let x1 = data[offset + 2 * i + 1];
    506
    8
                    let cos_val = cos_vals[i];
    507
    8
                    let sin_val = sin_vals[i];
    508
    8
    509
    8
                    // Apply 2D rotation with mscale
    510
    8
                    let y0 = (x0 * cos_val - x1 * sin_val) * self.mscale;
    511
    8
                    let y1 = (x0 * sin_val + x1 * cos_val) * self.mscale;
    512
    8
    513
    8
                    output.push(y0);
    514
    8
                    output.push(y1);
    515
    8
                }
    516
            }
    517
    518
    4
            Tensor::from_vec(shape.to_vec(), output)
    519
    5
        }
    520
    521
        /// Get embedding dimension
    522
        #[must_use]
    523
    6
        pub fn dim(&self) -> usize {
    524
    6
            self.dim
    525
    6
        }
    526
    527
        /// Get original base frequency
    528
        #[must_use]
    529
    2
        pub fn original_base(&self) -> f32 {
    530
    2
            self.original_base
    531
    2
        }
    532
    533
        /// Get scaled base frequency
    534
        #[must_use]
    535
    6
        pub fn scaled_base(&self) -> f32 {
    536
    6
            self.scaled_base
    537
    6
        }
    538
    539
        /// Get scaling configuration
    540
        #[must_use]
    541
    0
        pub fn scaling(&self) -> &RopeScalingType {
    542
    0
            &self.scaling
    543
    0
        }
    544
    545
        /// Get inverse frequencies
    546
        #[must_use]
    547
    1
        pub fn inv_freq(&self) -> &[f32] {
    548
    1
            &self.inv_freq
    549
    1
        }
    550
    551
        /// Get attention scaling factor (mscale)
    552
        #[must_use]
    553
    5
        pub fn mscale(&self) -> f32 {
    554
    5
            self.mscale
    555
    5
        }
    556
    557
        /// Compute effective context length multiplier
    558
        ///
    559
        /// Returns the factor by which the original context length is extended
    560
        #[must_use]
    561
    5
        pub fn context_length_multiplier(&self) -> f32 {
    562
    5
            match &self.scaling {
    563
    1
                RopeScalingType::None => 1.0,
    564
    2
                RopeScalingType::Linear { scale1 } | RopeScalingType::Ntk { scale1 } => *scale,
    565
                RopeScalingType::DynamicNtk {
    566
    1
                    original_max_len,
    567
    1
                    target_max_len,
    568
                }
    569
                | RopeScalingType::Yarn {
    570
    1
                    original_max_len,
    571
    1
                    target_max_len,
    572
                    ..
    573
    2
                } => (*target_max_len as f32) / (*original_max_len as f32),
    574
            }
    575
    5
        }
    576
    }
    577
    578
    /// Attention with Linear Biases (`ALiBi`)
    579
    ///
    580
    /// Replaces traditional position embeddings by adding a static, non-learned
    581
    /// bias to query-key attention scores. The bias is proportional to the distance
    582
    /// between positions, with head-specific slopes that enable better length
    583
    /// extrapolation.
    584
    ///
    585
    /// # Algorithm
    586
    ///
    587
    /// For each attention head h, `ALiBi` adds the following bias to attention scores:
    588
    ///
    589
    /// ```text
    590
    /// bias[i, j] = -m[h] * |i - j|
    591
    /// ```
    592
    ///
    593
    /// where m[h] is the head-specific slope computed as:
    594
    /// - For powers of 2: m[h] = 2^(-8h/n) where n is the number of heads
    595
    /// - For non-powers of 2: interpolation between adjacent powers of 2
    596
    ///
    597
    /// # References
    598
    ///
    599
    /// - "Train Short, Test Long: Attention with Linear Biases Enables Input Length
    600
    ///   Extrapolation" - Press et al., ICLR 2022
    601
    /// - <https://arxiv.org/abs/2108.12409>
    602
    ///
    603
    /// # Example
    604
    ///
    605
    /// ```rust,ignore
    606
    /// use realizar::layers::ALiBi;
    607
    ///
    608
    /// // Create ALiBi for 8 attention heads
    609
    /// let alibi = ALiBi::new(8)?;
    610
    ///
    611
    /// // Get bias matrix for sequence length 10
    612
    /// let bias = alibi.get_bias(10)?;
    613
    ///
    614
    /// // Add to attention scores before softmax
    615
    /// // scores: [seq_len, seq_len, num_heads]
    616
    /// // scores = scores + bias
    617
    /// ```
    618
    #[derive(Debug, Clone)]
    619
    pub struct ALiBi {
    620
        /// Number of attention heads
    621
        num_heads: usize,
    622
        /// Head-specific slopes (one per head)
    623
        slopes: Vec<f32>,
    624
    }
    625
    626
    impl ALiBi {
    627
        /// Create a new `ALiBi` layer
    628
        ///
    629
        /// # Arguments
    630
        ///
    631
        /// * `num_heads` - Number of attention heads
    632
        ///
    633
        /// # Errors
    634
        ///
    635
        /// Returns error if `num_heads` is zero
    636
    20
        pub fn new(num_heads: usize) -> Result<Self> {
    637
    20
            if num_heads == 0 {
    638
    2
                return Err(RealizarError::InvalidShape {
    639
    2
                    reason: "num_heads must be > 0".to_string(),
    640
    2
                });
    641
    18
            }
    642
    643
            // Compute slopes for each head
    644
    18
            let slopes = Self::compute_slopes(num_heads);
    645
    646
    18
            Ok(Self { num_heads, slopes })
    647
    20
        }
    648
    649
        /// Compute head-specific slopes following `ALiBi` paper algorithm
    650
        ///
    651
        /// For powers of 2: m[h] = 2^(-8h/n)
    652
        /// For non-powers of 2: interpolate between adjacent powers of 2
    653
    18
        fn compute_slopes(num_heads: usize) -> Vec<f32> {
    654
            // Find closest power of 2
    655
    18
            let closest_power_of_2 = if num_heads.is_power_of_two() {
    656
    15
                num_heads
    657
            } else {
    658
    3
                num_heads.next_power_of_two() / 2
    659
            };
    660
    661
            #[allow(clippy::cast_precision_loss)]
    662
    18
            let ratio = 8.0 / (closest_power_of_2 as f32);
    663
    664
    18
            let mut slopes = Vec::with_capacity(num_heads);
    665
    666
            // Compute slopes for power of 2 heads
    667
    102
            for i in 0..closest_power_of_218.min18(num_heads18) {
    668
    102
                #[allow(clippy::cast_precision_loss)]
    669
    102
                let exponent = -(i as f32) * ratio;
    670
    102
                slopes.push(2_f32.powf(exponent));
    671
    102
            }
    672
    673
            // If not power of 2, add extra slopes with step=2
    674
    18
            if num_heads > closest_power_of_2 {
    675
                #[allow(clippy::cast_precision_loss)]
    676
    3
                let extra_ratio = 4.0 / (closest_power_of_2 as f32);
    677
    678
    10
                for i in 0..(num_heads - closest_power_of_2)3 {
    679
    10
                    #[allow(clippy::cast_precision_loss)]
    680
    10
                    let exponent = -((2 * i + 1) as f32) * extra_ratio;
    681
    10
                    slopes.push(2_f32.powf(exponent));
    682
    10
                }
    683
    15
            }
    684
    685
    18
            slopes
    686
    18
        }
    687
    688
        /// Get bias matrix for a given sequence length
    689
        ///
    690
        /// Returns a tensor of shape `[seq_len, seq_len, num_heads]` where:
    691
        /// ```text
    692
        /// bias[i, j, h] = -slopes[h] * abs(i - j)
    693
        /// ```
    694
        ///
    695
        /// # Arguments
    696
        ///
    697
        /// * `seq_len` - Sequence length for computing bias
    698
        ///
    699
        /// # Returns
    700
        ///
    701
        /// Tensor of shape `[seq_len, seq_len, num_heads]` containing position biases
    702
        ///
    703
        /// # Errors
    704
        ///
    705
        /// Returns error if `seq_len` is zero
    706
    9
        pub fn get_bias(&self, seq_len: usize) -> Result<Tensor<f32>> {
    707
    9
            if seq_len == 0 {
    708
    1
                return Err(RealizarError::InvalidShape {
    709
    1
                    reason: "seq_len must be > 0".to_string(),
    710
    1
                });
    711
    8
            }
    712
    713
    8
            let total_size = seq_len * seq_len * self.num_heads;
    714
    8
            let mut data = Vec::with_capacity(total_size);
    715
    716
            // Compute bias for each position pair and head
    717
    173
            for i in 0..seq_len8 {
    718
    16.7k
                for j in 0..seq_len173 {
    719
    149k
                    for &slope132k in &self.slopes {
    720
    132k
                        #[allow(clippy::cast_precision_loss)]
    721
    132k
                        let distance = (i as f32 - j as f32).abs();
    722
    132k
                        let bias = -slope * distance;
    723
    132k
                        data.push(bias);
    724
    132k
                    }
    725
                }
    726
            }
    727
    728
    8
            Tensor::from_vec(vec![seq_len, seq_len, self.num_heads], data)
    729
    9
        }
    730
    731
        /// Get number of attention heads
    732
        #[must_use]
    733
    8
        pub fn num_heads(&self) -> usize {
    734
    8
            self.num_heads
    735
    8
        }
    736
    737
        /// Get head-specific slopes
    738
        #[must_use]
    739
    9
        pub fn slopes(&self) -> &[f32] {
    740
    9
            &self.slopes
    741
    9
        }
    742
    }