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
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

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

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

use super::{Embedding, Linear};

use crate::ml::{Module, ModuleRegistry, Parameter, byte};

// OA donor: oacpp/source/cpp/lib/oa/ml/nn/byte.cpp
// Adaptation: Complete existing Rust composition over canonical Embedding and Linear; preserves fixed 256-byte vocabulary, child-owned parameter identities, U8/U32 gather adjoints and FP32 affine adjoints; retains Rust deterministic seeded initialization rather than donor constructor distributions
/// Fixed 256-row trainable embedding for packed U8 or U32 byte IDs.
pub struct ByteEmbedding {
	inner: Embedding,
}

impl ByteEmbedding {
	/// Construct a deterministically initialized byte embedding.
	///
	/// # Errors
	///
	/// Returns an error when `embedding_dim` is zero or initialization fails.
	pub fn with_seed(engine: &Engine, embedding_dim: usize, seed: u64) -> Result<Self> {
		Ok(Self {
			inner: Embedding::with_seed(engine, byte::VOCAB_SIZE, embedding_dim, seed)?,
		})
	}

	/// Construct a byte embedding from an exact FP32 `[256, D]` table.
	///
	/// # Errors
	///
	/// Returns an error unless the table has exactly 256 nonempty rows.
	pub fn from_matrix(weight: Matrix) -> Result<Self> {
		if weight.shape().first() != Some(&byte::VOCAB_SIZE) {
			return Err(Error::invalid_argument(
				"byte embedding weight must have shape [256, D]",
			));
		}
		Ok(Self {
			inner: Embedding::from_matrix(weight)?,
		})
	}

	/// Gather one row per packed U8 or U32 byte ID.
	///
	/// # Errors
	///
	/// Returns an error for an invalid input or runtime recording failure.
	pub fn forward(&self, byte_ids: &Matrix) -> Result<Matrix> {
		self.inner.forward(byte_ids)
	}

	/// Embedding width.
	pub const fn embedding_dim(&self) -> usize {
		self.inner.embedding_dim()
	}

	/// Stable trainable table handle.
	pub fn weight(&self) -> Parameter {
		self.inner.weight()
	}
}

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

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

// OA donor: oacpp/source/cpp/lib/oa/ml/nn/byte.cpp
// Adaptation: Complete existing Rust composition over canonical Embedding and Linear; preserves fixed 256-byte vocabulary, child-owned parameter identities, U8/U32 gather adjoints and FP32 affine adjoints; retains Rust deterministic seeded initialization rather than donor constructor distributions
/// Trainable affine output head from hidden rows to 256 byte logits.
pub struct ByteHead {
	inner: Linear,
}

impl ByteHead {
	/// Construct a deterministically initialized byte head.
	///
	/// # Errors
	///
	/// Returns an error when `input_features` is zero or initialization fails.
	pub fn with_seed(engine: &Engine, input_features: usize, seed: u64) -> Result<Self> {
		Ok(Self {
			inner: Linear::with_seed(engine, input_features, byte::VOCAB_SIZE, seed)?,
		})
	}

	/// Project hidden states to 256 byte logits.
	///
	/// # Errors
	///
	/// Returns an error for an invalid input or runtime recording failure.
	pub fn forward(&self, hidden: &Matrix) -> Result<Matrix> {
		self.inner.forward(hidden)
	}

	/// Hidden input width.
	pub const fn input_features(&self) -> usize {
		self.inner.input_features()
	}

	/// Output weight handle.
	pub fn weight(&self) -> Parameter {
		self.inner.weight()
	}

	/// Output bias handle.
	pub fn bias(&self) -> Parameter {
		self
			.inner
			.bias()
			.expect("byte head always constructs a trainable bias")
	}
}

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

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