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

use crate::{Matrix, Result};

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

pub(in crate::ml) fn record_flow_linear_match(
	clean: &Matrix,
	noise: &Matrix,
	time: &Matrix,
	state: &Matrix,
	velocity: &Matrix,
) -> Result<()> {
	record_node(GradNode::FlowLinearState {
		clean: clean.clone(),
		noise: noise.clone(),
		time: time.clone(),
		output_id: state.value_id(),
	})?;
	record_node(GradNode::FlowLinearVelocity {
		clean: clean.clone(),
		noise: noise.clone(),
		output_id: velocity.value_id(),
	})
}

pub(in crate::ml) fn record_flow_euler_step(
	state: &Matrix,
	velocity: &Matrix,
	delta_time: f32,
	output: &Matrix,
) -> Result<()> {
	record_node(GradNode::FlowEulerStep {
		state: state.clone(),
		velocity: velocity.clone(),
		delta_time,
		output_id: output.value_id(),
	})
}

pub(in crate::ml) fn record_flow_masked_mse(
	prediction: &Matrix,
	target: &Matrix,
	mask: &Matrix,
	denominator: &Matrix,
	output: &Matrix,
) -> Result<()> {
	record_node(GradNode::FlowMaskedMse {
		prediction: prediction.clone(),
		target: target.clone(),
		mask: mask.clone(),
		denominator: denominator.clone(),
		output_id: output.value_id(),
	})
}