1use neura_graph::{Graph, Shape, Value};
2
3pub fn rms(graph: &Graph, value: Value, weight: Value, floor: Value) -> Value {
4 let squares = graph.mul(value, value);
5 let mean = graph.mean_axis(squares, 3);
6 let scale = graph.recip(graph.sqrt(graph.add(mean, floor)));
7 graph.mul(graph.mul(value, scale), weight)
8}
9
10pub fn l2(graph: &Graph, value: Value, floor: Value) -> Value {
11 let squares = graph.mul(value, value);
12 let total = graph.sum_axis(squares, 3);
13 let scale = graph.recip(graph.sqrt(graph.add(total, floor)));
14 graph.mul(value, scale)
15}
16
17pub fn softplus(graph: &Graph, value: Value) -> Value {
18 let threshold = graph.fill(Shape::scalar(), 20.0);
19 let raised = graph.log(graph.add(graph.fill(Shape::scalar(), 1.0), graph.exp(value)));
20 graph.select(graph.greater(value, threshold), value, raised)
21}
22
23pub fn joined(graph: &Graph, value: Value, count: u32) -> Value {
24 let shape = graph.shape(value);
25 assert_eq!(
26 shape.dims()[0],
27 count,
28 "a join lays the {count} heads of a chunk beside each other, and {:?} holds {} planes",
29 shape.dims(),
30 shape.dims()[0],
31 );
32 let parts = (0..count)
33 .map(|head| graph.slice(value, 0, head, 1))
34 .collect::<Vec<Value>>();
35 graph.concat(&parts, 3)
36}