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

// @generated by tool/gen/fn/generate.py; DO NOT EDIT.
// schema=ml_training.json schema_version=1 generator_version=64
// schema_sha256=6678961ee57f24310158676c7b4fe0821e21ca48db7edcb20d2a749a164976ae
// Saved-state and backward recipes; lifecycle remains tape-owned.
use crate::ml::autograd::node::operation::BackwardStatus;
use crate::ml::autograd::tape::backward::accumulate_value_gradient;
use crate::ml::matrix as ml_matrix;
use crate::{Matrix, Result};
use std::collections::HashMap;

pub(in crate::ml::autograd) struct GradNodeSwiglu {
	pub(in crate::ml::autograd) gate: Matrix,
	pub(in crate::ml::autograd) up: Matrix,
	pub(in crate::ml::autograd) output_id: u64,
}

impl GradNodeSwiglu {
	pub(in crate::ml::autograd) fn backward(
		&self,
		gradients: &mut HashMap<u64, Matrix>,
	) -> Result<BackwardStatus> {
		let Self {
			gate,
			up,
			output_id,
		} = self;
		let Some(output_gradient) = gradients.remove(output_id) else {
			return Ok(BackwardStatus::Skipped);
		};
		let (gate_gradient, up_gradient) = ml_matrix::swiglu_backward(gate, up, &output_gradient)?;
		accumulate_value_gradient(gradients, gate.value_id(), gate_gradient)?;
		accumulate_value_gradient(gradients, up.value_id(), up_gradient)?;
		Ok(BackwardStatus::Recorded)
	}
}

pub(in crate::ml::autograd) struct GradNodeSiluMul {
	pub(in crate::ml::autograd) input: Matrix,
	pub(in crate::ml::autograd) intermediate_size: usize,
	pub(in crate::ml::autograd) output_id: u64,
}

impl GradNodeSiluMul {
	pub(in crate::ml::autograd) fn backward(
		&self,
		gradients: &mut HashMap<u64, Matrix>,
	) -> Result<BackwardStatus> {
		let Self {
			input,
			intermediate_size,
			output_id,
		} = self;
		let Some(output_gradient) = gradients.remove(output_id) else {
			return Ok(BackwardStatus::Skipped);
		};
		let input_gradient =
			crate::ml::matrix::silu_mul_backward(input, &output_gradient, *intermediate_size)?;
		accumulate_value_gradient(gradients, input.value_id(), input_gradient)?;
		Ok(BackwardStatus::Recorded)
	}
}