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

// @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::{Matrix, Result};
use std::collections::HashMap;

pub(in crate::ml::autograd) struct GradNodeAvgPool2d {
	pub(in crate::ml::autograd) input: Matrix,
	pub(in crate::ml::autograd) kernel_size: usize,
	pub(in crate::ml::autograd) stride: usize,
	pub(in crate::ml::autograd) padding: usize,
	pub(in crate::ml::autograd) output_id: u64,
}

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

pub(in crate::ml::autograd) struct GradNodeMaxPool2d {
	pub(in crate::ml::autograd) input: Matrix,
	pub(in crate::ml::autograd) indices: Matrix,
	pub(in crate::ml::autograd) kernel_size: usize,
	pub(in crate::ml::autograd) stride: usize,
	pub(in crate::ml::autograd) padding: usize,
	pub(in crate::ml::autograd) output_id: u64,
}

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

pub(in crate::ml::autograd) struct GradNodeAdaptiveAvgPool2d {
	pub(in crate::ml::autograd) input: Matrix,
	pub(in crate::ml::autograd) output_height: usize,
	pub(in crate::ml::autograd) output_width: usize,
	pub(in crate::ml::autograd) output_id: u64,
}

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