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
104
105
106
107
108
109
110
111
112
113
114

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

use std::rc::Rc;

use crate::{Engine, Matrix, Result, matrix as core_matrix};

use super::super::{Module, ModuleRegistry, matrix::swiglu};

use super::{Linear, RmsNorm};

// OA donor: oacpp/source/cpp/lib/oa/ml/nn/ffn/ffn.cpp
// Adaptation: Complete generation of existing Rust donor composition; preserved Rc child ownership, seeded Xavier projections and seven parameter paths; Matrix and child modules retain adjoints
/// Pre-normalized SwiGLU feed-forward network with a residual connection.
///
/// The module preserves the OA donor structure:
/// `RMSNorm -> (gate, up) -> SwiGLU -> down -> residual add`.
pub struct Ffn {
	model_width: usize,
	hidden_width: usize,
	epsilon: f32,
	norm: Rc<RmsNorm>,
	gate: Rc<Linear>,
	up: Rc<Linear>,
	down: Rc<Linear>,
	registry: ModuleRegistry,
}

impl Ffn {
	/// Construct a deterministically initialized donor-compatible FFN.
	///
	/// # Errors
	///
	/// Returns an error for invalid dimensions or epsilon, or when child
	/// construction or registration fails.
	pub fn with_seed(
		engine: &Engine,
		model_width: usize,
		hidden_width: usize,
		epsilon: f32,
		seed: u64,
	) -> Result<Self> {
		let norm = Rc::new(RmsNorm::new(engine, model_width, epsilon)?);
		let gate = Rc::new(Linear::with_seed(engine, model_width, hidden_width, seed)?);
		let up = Rc::new(Linear::with_seed(
			engine,
			model_width,
			hidden_width,
			seed.wrapping_add(1),
		)?);
		let down = Rc::new(Linear::with_seed(
			engine,
			hidden_width,
			model_width,
			seed.wrapping_add(2),
		)?);
		let mut registry = ModuleRegistry::new();
		registry.register_module("norm", norm.clone())?;
		registry.register_module("gate", gate.clone())?;
		registry.register_module("up", up.clone())?;
		registry.register_module("down", down.clone())?;
		Ok(Self {
			model_width,
			hidden_width,
			epsilon,
			norm,
			gate,
			up,
			down,
			registry,
		})
	}

	/// Apply the pre-normalized gated projection and residual connection.
	///
	/// # Errors
	///
	/// Returns an error unless the input is a compatible F32 Matrix or a child
	/// operation fails.
	pub fn forward(&self, input: &Matrix) -> Result<Matrix> {
		let normalized = self.norm.forward(input)?;
		let gate = self.gate.forward(&normalized)?;
		let up = self.up.forward(&normalized)?;
		let activated = swiglu(&gate, &up)?;
		let projected = self.down.forward(&activated)?;
		core_matrix::add(input, &projected)
	}

	/// Return the input and output feature width.
	pub const fn model_width(&self) -> usize {
		self.model_width
	}

	/// Return the gated hidden feature width.
	pub const fn hidden_width(&self) -> usize {
		self.hidden_width
	}

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

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

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