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_softmax.json
// generator_version=8; schema_sha256=d733db02905d6952e4a537742dd04991a4596a5f7fad51b723228a35c21c3282

use crate::{Matrix, Result, matrix};

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

// OA donor: oacpp/source/cpp/lib/oa/ml/nn/softmax/softmax.cpp
// Adaptation: Rust ownership, ModuleRegistry and checked error contracts; preserved operation behavior
/// Parameterless axis-aware Softmax module.
pub struct Softmax {
	dim: i32,
	registry: ModuleRegistry,
}

impl Softmax {
	/// Construct a Softmax over `dim`, where `-1` selects the last dimension.
	pub const fn new(dim: i32) -> Self {
		Self {
			dim,
			registry: ModuleRegistry::new(),
		}
	}

	/// Return the selected dimension.
	pub const fn dim(&self) -> i32 {
		self.dim
	}

	/// Apply Softmax without submitting or waiting.
	///
	/// # Errors
	///
	/// Returns an error from the underlying Matrix operation.
	pub fn forward(&self, input: &Matrix) -> Result<Matrix> {
		matrix::softmax(input, self.dim)
	}
}

impl Default for Softmax {
	fn default() -> Self {
		Self::new(-1)
	}
}

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

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

// OA donor: oacpp/source/cpp/lib/oa/ml/nn/softmax/softmax.cpp
// Adaptation: Rust ownership, ModuleRegistry and checked error contracts; preserved operation behavior
/// Parameterless axis-aware LogSoftmax module.
pub struct LogSoftmax {
	dim: i32,
	registry: ModuleRegistry,
}

impl LogSoftmax {
	/// Construct a LogSoftmax over `dim`, where `-1` selects the last dimension.
	pub const fn new(dim: i32) -> Self {
		Self {
			dim,
			registry: ModuleRegistry::new(),
		}
	}

	/// Return the selected dimension.
	pub const fn dim(&self) -> i32 {
		self.dim
	}

	/// Apply LogSoftmax without submitting or waiting.
	///
	/// # Errors
	///
	/// Returns an error from the underlying Matrix operation.
	pub fn forward(&self, input: &Matrix) -> Result<Matrix> {
		matrix::log_softmax(input, self.dim)
	}
}

impl Default for LogSoftmax {
	fn default() -> Self {
		Self::new(-1)
	}
}

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

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