Skip to main content

deep_delta_learning/
utils.rs

1use burn::prelude::*;
2use burn::tensor::Int;
3
4pub fn tensor_to_vec<const D: usize, B: Backend>(tensor: Tensor<B, D>) -> Vec<f32> {
5    tensor
6        .into_data()
7        .to_vec::<f32>()
8        .expect("tensor conversion to f32 vec should succeed")
9}
10
11pub fn int_tensor_to_vec<const D: usize, B: Backend>(tensor: Tensor<B, D, Int>) -> Vec<i64> {
12    tensor
13        .into_data()
14        .to_vec::<i64>()
15        .expect("tensor conversion to i64 vec should succeed")
16}
17
18pub fn mean_of_slice(values: &[f32]) -> f32 {
19    if values.is_empty() {
20        return 0.0;
21    }
22    values.iter().copied().sum::<f32>() / values.len() as f32
23}
24
25pub fn std_of_slice(values: &[f32]) -> f32 {
26    if values.is_empty() {
27        return 0.0;
28    }
29    let mean = mean_of_slice(values);
30    let variance = values
31        .iter()
32        .map(|value| {
33            let centered = *value - mean;
34            centered * centered
35        })
36        .sum::<f32>()
37        / values.len() as f32;
38    variance.sqrt()
39}
40
41pub fn cosine_similarity(left: &[f32], right: &[f32]) -> f32 {
42    if left.len() != right.len() || left.is_empty() {
43        return 0.0;
44    }
45
46    let dot = left
47        .iter()
48        .zip(right.iter())
49        .map(|(lhs, rhs)| lhs * rhs)
50        .sum::<f32>();
51    let left_norm = left.iter().map(|value| value * value).sum::<f32>().sqrt();
52    let right_norm = right.iter().map(|value| value * value).sum::<f32>().sqrt();
53
54    if left_norm == 0.0 || right_norm == 0.0 {
55        0.0
56    } else {
57        dot / (left_norm * right_norm)
58    }
59}
60
61pub fn tensor_mean<const D: usize, B: Backend>(tensor: Tensor<B, D>) -> f32 {
62    mean_of_slice(&tensor_to_vec(tensor))
63}
64
65pub fn tensor_l2_norm<const D: usize, B: Backend>(tensor: Tensor<B, D>) -> f32 {
66    let values = tensor_to_vec(tensor);
67    values.iter().map(|value| value * value).sum::<f32>().sqrt()
68}
69
70pub fn create_causal_mask<B: Backend>(
71    batch_size: usize,
72    seq_len: usize,
73    device: &B::Device,
74) -> Tensor<B, 3> {
75    Tensor::<B, 2>::ones([seq_len, seq_len], device)
76        .triu(1)
77        .mul_scalar(-1.0e9f32)
78        .unsqueeze_dim::<3>(0)
79        .repeat_dim(0, batch_size)
80}
81
82pub fn beta_bias(beta_init: f64) -> f64 {
83    let probability = beta_init / 2.0;
84    (probability / (1.0 - probability)).ln()
85}
86
87pub fn mean_last_dim<const D: usize, B: Backend>(tensor: Tensor<B, D>) -> Tensor<B, D> {
88    tensor.mean_dim(D - 1)
89}
90
91pub fn scalar_as_tensor<B: Backend, const D: usize>(
92    value: f32,
93    shape: [usize; D],
94    device: &B::Device,
95) -> Tensor<B, D> {
96    Tensor::full(shape, value, device)
97}