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

// @generated by tool/gen/fn/generate.py; DO NOT EDIT.
// schema=ml_training.json schema_version=1 generator_version=64
// schema_sha256=6678961ee57f24310158676c7b4fe0821e21ca48db7edcb20d2a749a164976ae
// Mechanical tape attachments; concrete saved-state node behavior remains private.

use crate::{Matrix, Result};

use crate::ml::autograd::{
	node::{GradNode, operation::GradNodeOperation},
	tape::record_node,
};

use crate::ml::autograd::node::operation::{
	GradNodeBce, GradNodeCrossEntropy, GradNodeL1, GradNodeMaskedCrossEntropy, GradNodeMse,
	GradNodePpoClippedPolicy, GradNodeSmoothL1,
};

pub(in crate::ml) fn record_ppo_clipped_policy(
	new_log_probability: &Matrix,
	old_log_probability: &Matrix,
	advantage: &Matrix,
	clip_epsilon: f32,
	output: &Matrix,
) -> Result<()> {
	record_node(GradNode::Operation(GradNodeOperation::PpoClippedPolicy(
		GradNodePpoClippedPolicy {
			new_log_probability: new_log_probability.clone(),
			old_log_probability: old_log_probability.clone(),
			advantage: advantage.clone(),
			clip_epsilon,
			output_id: output.value_id(),
		},
	)))
}

pub(in crate::ml) fn record_cross_entropy(
	logits: &Matrix,
	targets: &Matrix,
	result: &Matrix,
) -> Result<()> {
	record_node(GradNode::Operation(GradNodeOperation::CrossEntropy(
		GradNodeCrossEntropy {
			logits: logits.clone(),
			targets: targets.clone(),
			output_id: result.value_id(),
		},
	)))
}

pub(in crate::ml) fn record_masked_cross_entropy(
	logits: &Matrix,
	targets: &Matrix,
	mask: &Matrix,
	valid_count: usize,
	result: &Matrix,
) -> Result<()> {
	record_node(GradNode::Operation(GradNodeOperation::MaskedCrossEntropy(
		GradNodeMaskedCrossEntropy {
			logits: logits.clone(),
			targets: targets.clone(),
			mask: mask.clone(),
			valid_count,
			output_id: result.value_id(),
		},
	)))
}

pub(in crate::ml) fn record_smooth_l1(
	prediction: &Matrix,
	target: &Matrix,
	result: &Matrix,
) -> Result<()> {
	record_node(GradNode::Operation(GradNodeOperation::SmoothL1(
		GradNodeSmoothL1 {
			prediction: prediction.clone(),
			target: target.clone(),
			output_id: result.value_id(),
		},
	)))
}

pub(in crate::ml) fn record_mse(
	prediction: &Matrix,
	target: &Matrix,
	result: &Matrix,
) -> Result<()> {
	record_node(GradNode::Operation(GradNodeOperation::Mse(GradNodeMse {
		prediction: prediction.clone(),
		target: target.clone(),
		output_id: result.value_id(),
	})))
}

pub(in crate::ml) fn record_l1(
	prediction: &Matrix,
	target: &Matrix,
	result: &Matrix,
) -> Result<()> {
	record_node(GradNode::Operation(GradNodeOperation::L1(GradNodeL1 {
		prediction: prediction.clone(),
		target: target.clone(),
		output_id: result.value_id(),
	})))
}

pub(in crate::ml) fn record_bce(
	prediction: &Matrix,
	target: &Matrix,
	result: &Matrix,
) -> Result<()> {
	record_node(GradNode::Operation(GradNodeOperation::Bce(GradNodeBce {
		prediction: prediction.clone(),
		target: target.clone(),
		output_id: result.value_id(),
	})))
}