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

use crate::{Matrix, Result};

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

pub(in crate::ml) fn record_bmm(left: &Matrix, right: &Matrix, output: &Matrix) -> Result<()> {
	record_node(GradNode::Bmm {
		left: left.clone(),
		right: right.clone(),
		output_id: output.value_id(),
	})
}

pub(in crate::ml) fn record_bmm_nt(left: &Matrix, right: &Matrix, output: &Matrix) -> Result<()> {
	record_node(GradNode::BmmNt {
		left: left.clone(),
		right: right.clone(),
		output_id: output.value_id(),
	})
}

pub(in crate::ml) fn record_bmm_tn(left: &Matrix, right: &Matrix, output: &Matrix) -> Result<()> {
	record_node(GradNode::BmmTn {
		left: left.clone(),
		right: right.clone(),
		output_id: output.value_id(),
	})
}

pub(in crate::ml) fn record_split_heads(
	input: &Matrix,
	output: &Matrix,
	batch: usize,
	sequence_length: usize,
	num_heads: usize,
) -> Result<()> {
	record_node(GradNode::SplitHeads {
		input: input.clone(),
		output_id: output.value_id(),
		batch,
		sequence_length,
		num_heads,
	})
}

pub(in crate::ml) fn record_merge_heads(
	input: &Matrix,
	output: &Matrix,
	batch: usize,
	sequence_length: usize,
	num_heads: usize,
) -> Result<()> {
	record_node(GradNode::MergeHeads {
		input: input.clone(),
		output_id: output.value_id(),
		batch,
		sequence_length,
		num_heads,
	})
}

pub(in crate::ml) fn record_softmax_scaled_masked(
	input: &Matrix,
	output: &Matrix,
	scale: f32,
) -> Result<()> {
	record_node(GradNode::SoftmaxScaledMasked {
		input: input.clone(),
		output: output.clone(),
		output_id: output.value_id(),
		scale,
	})
}

pub(in crate::ml) fn record_scaled_dot_product_attention(
	query: &Matrix,
	key: &Matrix,
	value: &Matrix,
	output: &Matrix,
	probabilities: Matrix,
	scale: f32,
) -> Result<()> {
	record_node(GradNode::ScaledDotProductAttention {
		query: query.clone(),
		key: key.clone(),
		value: value.clone(),
		probabilities,
		output_id: output.value_id(),
		scale,
	})
}

pub(in crate::ml) fn record_flash_attention(
	query: &Matrix,
	key: &Matrix,
	value: &Matrix,
	output: &Matrix,
	log_sum_exp: Matrix,
	scale: f32,
) -> Result<()> {
	record_node(GradNode::FlashAttention {
		query: query.clone(),
		key: key.clone(),
		value: value.clone(),
		output: output.clone(),
		log_sum_exp,
		output_id: output.value_id(),
		scale,
	})
}