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

use crate::ml::matrix::{Mamba3PreprocessConfig, Mamba3PreprocessResult, SsmConfig};
use crate::{Matrix, Result};

use super::super::{node::GradNode, tape::record_node};

pub(in crate::ml) fn record_mamba3_preprocess(
	projected: &Matrix,
	dt_bias: &Matrix,
	outputs: &Mamba3PreprocessResult,
	config: Mamba3PreprocessConfig,
) -> Result<()> {
	let [x, z, bh, ch, dt, adt, trap, angle] = outputs.matrices();
	record_node(GradNode::Mamba3Preprocess(Box::new(
		super::super::node::GradNodeMamba3Preprocess {
			projected: projected.clone(),
			dt_bias: dt_bias.clone(),
			outputs: [
				x.clone(),
				z.clone(),
				bh.clone(),
				ch.clone(),
				dt.clone(),
				adt.clone(),
				trap.clone(),
				angle.clone(),
			],
			config,
		},
	)))
}

pub(in crate::ml) fn record_mamba3_siso(
	inputs: [&Matrix; 11],
	output: &Matrix,
	config: SsmConfig,
) -> Result<()> {
	let [c, b, x, z, adt, dt, trap, angle, c_bias, b_bias, d] = inputs;
	record_node(GradNode::Mamba3Siso(Box::new(
		super::super::node::GradNodeMamba3Siso {
			c: c.clone(),
			b: b.clone(),
			x: x.clone(),
			z: z.clone(),
			adt: adt.clone(),
			dt: dt.clone(),
			trap: trap.clone(),
			angle: angle.clone(),
			c_bias: c_bias.clone(),
			b_bias: b_bias.clone(),
			d: d.clone(),
			config,
			output_id: output.value_id(),
		},
	)))
}

pub(in crate::ml) fn record_mamba3_mimo(
	inputs: [&Matrix; 15],
	output: &Matrix,
	config: SsmConfig,
) -> Result<()> {
	record_node(GradNode::Mamba3Mimo(Box::new(
		super::super::node::GradNodeMamba3Mimo {
			inputs: inputs.map(Clone::clone),
			config,
			output_id: output.value_id(),
		},
	)))
}