1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
// @generated by tool/gen/nn/generate.py; DO NOT EDIT.
// Source: tool/gen/nn/schema/ml/ml_nn_rope.json
// generator_version=8; schema_sha256=ad31cadd74a43cd80eff08346285c2feba4708a65ac92fae380b49669120d64e
use crate::{Error, Matrix, Result};
use super::super::{Module, ModuleRegistry, matrix};
// OA donor: oacpp/source/cpp/lib/oa/ml/nn/rope/rope.cpp
// Adaptation: Preserved checked head geometry, finite positive theta and zero position offset; existing Matrix provider owns donor split-half rotation and adjoint
/// Llama-style split-half rotary position embedding module.
pub struct Rope {
num_heads: usize,
head_dim: usize,
theta_base: f32,
registry: ModuleRegistry,
}
impl Rope {
/// Construct a parameterless rotary-position module.
///
/// # Errors
///
/// Returns an error unless both dimensions are nonzero, `head_dim` is even,
/// their product fits `usize`, and `theta_base` is finite and positive.
pub fn new(num_heads: usize, head_dim: usize, theta_base: f32) -> Result<Self> {
if num_heads == 0
|| head_dim == 0
|| !head_dim.is_multiple_of(2)
|| !theta_base.is_finite()
|| theta_base <= 0.0
|| num_heads.checked_mul(head_dim).is_none()
{
return Err(Error::invalid_argument(
"RoPE requires nonzero heads, an even head dimension, and a finite positive theta base",
));
}
Ok(Self {
num_heads,
head_dim,
theta_base,
registry: ModuleRegistry::new(),
})
}
/// Rotate a sequence beginning at absolute position zero.
///
/// # Errors
///
/// Returns an error unless input is compatible F32 `[T, H * D]`, or runtime
/// recording fails.
pub fn forward(&self, input: &Matrix) -> Result<Matrix> {
matrix::rope(input, self.num_heads, self.head_dim, self.theta_base, 0)
}
/// Return the number of independently rotated heads.
pub const fn num_heads(&self) -> usize {
self.num_heads
}
/// Return the even dimension of each head.
pub const fn head_dim(&self) -> usize {
self.head_dim
}
/// Return the rotary frequency base.
pub const fn theta_base(&self) -> f32 {
self.theta_base
}
}
impl Module for Rope {
fn forward(&self, input: &Matrix) -> Result<Matrix> {
Rope::forward(self, input)
}
fn registry(&self) -> &ModuleRegistry {
&self.registry
}
}