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

use crate::{Matrix, Result, ml::Parameter};

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

pub(in crate::ml) fn record_grouped_gemm_m(
	input: &Matrix,
	output: &Matrix,
	weight: Option<(Parameter, u64)>,
	weight_value: Matrix,
	offsets: &Matrix,
) -> Result<()> {
	let (weight, weight_version) = weight
		.map(|(parameter, version)| (Some(parameter), Some(version)))
		.unwrap_or((None, None));
	record_node(GradNode::GroupedGemmM {
		input: input.clone(),
		weight,
		weight_value,
		weight_version,
		offsets: offsets.clone(),
		output_id: output.value_id(),
	})
}

#[allow(
	clippy::too_many_arguments,
	reason = "the node retains exact grouped-Linear values and optional module parameters"
)]
pub(in crate::ml) fn record_grouped_linear_m(
	input: &Matrix,
	output: &Matrix,
	weight: Option<(Parameter, u64)>,
	weight_value: Matrix,
	bias: Option<(Parameter, u64)>,
	bias_value: Matrix,
	offsets: &Matrix,
) -> Result<()> {
	let (weight, weight_version) = weight
		.map(|(parameter, version)| (Some(parameter), Some(version)))
		.unwrap_or((None, None));
	let (bias, bias_version) = bias
		.map(|(parameter, version)| (Some(parameter), Some(version)))
		.unwrap_or((None, None));
	record_node(GradNode::GroupedLinearM {
		input: input.clone(),
		weight,
		weight_value,
		weight_version,
		bias,
		bias_value,
		bias_version,
		offsets: offsets.clone(),
		output_id: output.value_id(),
	})
}

pub(in crate::ml) fn record_moe_route_weights(
	probabilities: &Matrix,
	expert_indices: &Matrix,
	output: &Matrix,
) -> Result<()> {
	record_node(GradNode::MoeRouteWeights {
		probabilities: probabilities.clone(),
		expert_indices: expert_indices.clone(),
		output: output.clone(),
		output_id: output.value_id(),
	})
}

pub(in crate::ml) fn record_moe_gather(
	input: &Matrix,
	inverse: &Matrix,
	output: &Matrix,
) -> Result<()> {
	record_node(GradNode::MoeGather {
		input: input.clone(),
		inverse: inverse.clone(),
		output_id: output.value_id(),
	})
}

pub(in crate::ml) fn record_moe_combine(
	packed: &Matrix,
	route_gate: &Matrix,
	inverse: &Matrix,
	packed_slot: &Matrix,
	output: &Matrix,
) -> Result<()> {
	record_node(GradNode::MoeCombine {
		packed: packed.clone(),
		route_gate: route_gate.clone(),
		inverse: inverse.clone(),
		packed_slot: packed_slot.clone(),
		output_id: output.value_id(),
	})
}