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

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

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

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

// OA donor: oacpp/source/cpp/lib/oa/ml/nn/embedding/embedding.cpp
// Adaptation: Preserve existing Rust initialization, parameter versions and autograd; adapt donor structure to ModuleRegistry ownership
/// Trainable FP32 lookup table indexed by an arbitrary-shape U8, U32, or I32 matrix.
pub struct Embedding {
	num_embeddings: usize,
	embedding_dim: usize,
	weight: Parameter,
	registry: ModuleRegistry,
}

impl Embedding {
	/// Construct a deterministically initialized embedding table.
	///
	/// # Errors
	///
	/// Returns an error when either dimension is zero, shape arithmetic overflows,
	/// or allocation/upload fails.
	pub fn with_seed(
		engine: &Engine,
		num_embeddings: usize,
		embedding_dim: usize,
		seed: u64,
	) -> Result<Self> {
		if num_embeddings == 0 || embedding_dim == 0 {
			return Err(Error::invalid_argument(
				"embedding dimensions must be nonzero",
			));
		}
		let count = num_embeddings
			.checked_mul(embedding_dim)
			.ok_or_else(|| Error::invalid_argument("embedding weight size overflows usize"))?;
		let limit = 1.0_f32 / (embedding_dim as f32).sqrt();
		let values = random::symmetric_uniform(count, limit, seed);
		let weight = Matrix::from_f32(engine, [num_embeddings, embedding_dim], &values)?;
		Self::from_matrix(weight)
	}

	/// Construct an embedding from an exact FP32 table shaped `[V, D]`.
	///
	/// # Errors
	///
	/// Returns an error unless `weight` is a nonempty rank-two FP32 matrix.
	pub fn from_matrix(weight: Matrix) -> Result<Self> {
		let [num_embeddings, embedding_dim] = weight.shape() else {
			return Err(Error::invalid_argument(
				"embedding weight must have rank two",
			));
		};
		if *num_embeddings == 0 || *embedding_dim == 0 || weight.dtype() != DType::F32 {
			return Err(Error::invalid_argument(
				"embedding requires a nonempty FP32 weight [V, D]",
			));
		}
		let num_embeddings = *num_embeddings;
		let embedding_dim = *embedding_dim;
		let weight = Parameter::new("weight", weight)?;
		let mut registry = ModuleRegistry::new();
		registry.register_parameter("weight", weight.clone())?;
		Ok(Self {
			num_embeddings,
			embedding_dim,
			weight,
			registry,
		})
	}

	/// Gather one embedding row per U8, U32, or non-negative I32 index.
	///
	/// The output shape is the complete index shape followed by `embedding_dim`.
	/// Out-of-range indices produce NaN rows without reading outside the table.
	///
	/// # Errors
	///
	/// Returns an error when indices are not U8/U32/I32, the engines differ, shape
	/// or ABI arithmetic overflows, or runtime recording fails. Negative I32 and
	/// other out-of-range indices produce NaN rows without an out-of-bounds read.
	pub fn forward(&self, indices: &Matrix) -> Result<Matrix> {
		let (weight, version, requires_grad) = self.weight.snapshot();
		let output = matrix::embedding(&weight, indices)?;
		if requires_grad {
			autograd::record_embedding(indices, &output, self.weight.clone(), weight, version)?;
		}
		Ok(output)
	}

	pub const fn num_embeddings(&self) -> usize {
		self.num_embeddings
	}

	pub const fn embedding_dim(&self) -> usize {
		self.embedding_dim
	}

	pub fn weight(&self) -> Parameter {
		self.weight.clone()
	}

	pub fn parameters(&self) -> [Parameter; 1] {
		[self.weight.clone()]
	}
}

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

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