oarust 0.8.6

GPU-first semantic computing engine written in Rust
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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103

// @generated by tool/gen/nn/generate.py; DO NOT EDIT.
// Source: tool/gen/nn/schema/ml/ml_nn_rms_norm.json
// generator_version=8; schema_sha256=8c6285cdde7902e200d5e837b80438d4cd0443c5d7297a8bac84332f69593f58

use crate::{DType, Engine, Error, Matrix, Result};

use super::super::{Module, ModuleRegistry, Parameter, autograd, matrix};

// OA donor: oacpp/source/cpp/lib/oa/ml/nn/rmsnorm/rmsNorm.cpp
// Adaptation: Preserve existing Rust initialization, parameter versions and autograd; adapt donor structure to ModuleRegistry ownership
/// Trainable FP32 RMSNorm over the final input dimension.
pub struct RmsNorm {
	dimension: usize,
	epsilon: f32,
	weight: Parameter,
	registry: ModuleRegistry,
}

impl RmsNorm {
	/// Construct RMSNorm with unit weight.
	///
	/// # Errors
	///
	/// Returns an error when `dimension` is zero, `epsilon` is not finite and
	/// positive, or parameter allocation/upload fails.
	pub fn new(engine: &Engine, dimension: usize, epsilon: f32) -> Result<Self> {
		if dimension == 0 {
			return Err(Error::invalid_argument("RMSNorm dimension must be nonzero"));
		}
		let weight = Matrix::from_f32(engine, [dimension], &vec![1.0; dimension])?;
		Self::from_matrix(weight, epsilon)
	}

	/// Construct RMSNorm from an exact weight Matrix.
	///
	/// # Errors
	///
	/// Returns an error unless weight is a nonempty F32 vector and `epsilon` is
	/// finite and positive.
	pub fn from_matrix(weight: Matrix, epsilon: f32) -> Result<Self> {
		let [dimension] = weight.shape() else {
			return Err(Error::invalid_argument("RMSNorm weight must be rank one"));
		};
		if *dimension == 0 || weight.dtype() != DType::F32 || !epsilon.is_finite() || epsilon <= 0.0 {
			return Err(Error::invalid_argument(
				"RMSNorm requires a nonempty F32 weight vector and finite positive epsilon",
			));
		}
		let dimension = *dimension;
		let weight = Parameter::new("weight", weight)?;
		let mut registry = ModuleRegistry::new();
		registry.register_parameter("weight", weight.clone())?;
		Ok(Self {
			dimension,
			epsilon,
			weight,
			registry,
		})
	}

	/// Normalize each row by its root mean square and apply weight.
	///
	/// # Errors
	///
	/// Returns an error for an incompatible input or runtime recording failure.
	pub fn forward(&self, input: &Matrix) -> Result<Matrix> {
		let (weight, version, requires_grad) = self.weight.snapshot();
		let output = matrix::rms_norm_forward(input, &weight, self.epsilon)?;
		let parameter = requires_grad.then(|| (self.weight.clone(), version));
		autograd::record_rms_norm(input, &output, parameter, weight, self.epsilon)?;
		Ok(output)
	}

	/// Return the normalized final-dimension width.
	pub const fn dimension(&self) -> usize {
		self.dimension
	}

	/// Return the numerical-stability epsilon.
	pub const fn epsilon(&self) -> f32 {
		self.epsilon
	}

	/// Return the stable trainable weight handle.
	pub fn weight(&self) -> Parameter {
		self.weight.clone()
	}

	/// Return the one registered parameter.
	pub fn parameters(&self) -> [Parameter; 1] {
		[self.weight.clone()]
	}
}

impl Module for RmsNorm {
	fn forward(&self, input: &Matrix) -> Result<Matrix> {
		RmsNorm::forward(self, input)
	}

	fn registry(&self) -> &ModuleRegistry {
		&self.registry
	}
}