burn-tensor 0.18.0

Tensor library with user-friendly APIs and automatic differentiation support
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

#[burn_tensor_testgen::testgen(cosine_similarity)]
mod tests {
    use super::*;
    use burn_tensor::TensorData;
    use burn_tensor::linalg;
    use burn_tensor::{ElementConversion, Tolerance, ops::FloatElem};
    type FT = FloatElem<TestBackend>;

    #[test]
    fn test_cosine_similarity_basic() {
        // Create test tensors
        let x1 = TestTensor::<2>::from([[1.0, 2.0, 3.0], [0.5, 1.5, 2.5]]);
        let x2 = TestTensor::<2>::from([[1.5, 2.5, 3.5], [0.7, 1.7, 2.7]]);

        // Test cosine similarity along dimension 1
        let expected = TestTensor::<2>::from([[0.99983203], [0.99987257]]).into_data();
        linalg::cosine_similarity(x1.clone(), x2.clone(), 1, None)
            .into_data()
            .assert_approx_eq::<FT>(&expected, Tolerance::default());

        // Test with explicit epsilon
        linalg::cosine_similarity(x1.clone(), x2.clone(), 1, Some(1e-8.elem::<FT>()))
            .into_data()
            .assert_approx_eq::<FT>(&expected, Tolerance::default());
    }

    #[test]
    fn test_cosine_similarity_orthogonal() {
        // Create orthogonal vectors
        let x1 = TestTensor::<2>::from([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]);
        let x2 = TestTensor::<2>::from([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]);

        // Orthogonal vectors should have cosine similarity of 0
        let expected = TestTensor::<2>::from([[0.0], [0.0]]).into_data();
        linalg::cosine_similarity(x1, x2, 1, None)
            .into_data()
            .assert_eq(&expected, true);
    }

    #[test]
    fn test_cosine_similarity_parallel() {
        // Create parallel vectors
        let x1 = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);
        let x2 = TestTensor::<2>::from([[2.0, 4.0, 6.0], [8.0, 10.0, 12.0]]);

        // Parallel vectors should have cosine similarity of 1
        let expected = TestTensor::<2>::from([[1.0], [1.0]]).into_data();
        linalg::cosine_similarity(x1, x2, 1, None)
            .into_data()
            .assert_approx_eq::<FT>(&expected, Tolerance::default());
    }

    #[test]
    fn test_cosine_similarity_opposite() {
        // Create opposite direction vectors
        let x1 = TestTensor::<2>::from([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);
        let x2 = TestTensor::<2>::from([[-1.0, -2.0, -3.0], [-4.0, -5.0, -6.0]]);

        // Opposite vectors should have cosine similarity of -1
        let expected = TestTensor::<2>::from([[-1.0], [-1.0]]).into_data();
        linalg::cosine_similarity(x1, x2, 1, None)
            .into_data()
            .assert_approx_eq::<FT>(&expected, Tolerance::default());
    }

    #[test]
    fn test_cosine_similarity_different_dimension() {
        // Test with a 3D tensor
        let x1 = TestTensor::<3>::from([[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]);
        let x2 = TestTensor::<3>::from([[[2.0, 3.0], [4.0, 5.0]], [[6.0, 7.0], [8.0, 9.0]]]);

        // Test along dimension 2
        let expected =
            TestTensor::<3>::from([[[0.9959688], [0.9958376]], [[0.9955946], [0.9955169]]])
                .into_data();

        // sensitive to rounding in dot/norm; loosen f16 tolerance
        let tolerance = Tolerance::default().set_half_precision_relative(7e-3);

        linalg::cosine_similarity(x1.clone(), x2.clone(), 2, None)
            .into_data()
            .assert_approx_eq::<FT>(&expected, tolerance);

        // Test with negative dimension (-1 is the last dimension, which is 2 in this case)
        linalg::cosine_similarity(x1.clone(), x2.clone(), -1, None)
            .into_data()
            .assert_approx_eq::<FT>(&expected, tolerance);
    }

    #[test]
    fn test_cosine_similarity_near_zero() {
        // Test with near-zero vectors
        let x1 = TestTensor::<2>::from([[1e-10, 2e-10, 3e-10], [4e-10, 5e-10, 6e-10]]);
        let x2 = TestTensor::<2>::from([[2e-10, 4e-10, 6e-10], [8e-10, 10e-10, 12e-10]]);

        // Update the expected values based on the actual implementation behavior
        let expected = TestTensor::<2>::from([[0.0028], [0.0154]]).into_data();

        // Smaller values result in NaN on metal f16
        let epsilon = Some(FT::from_elem(1e-2));
        let tolerance = Tolerance::absolute(0.2);

        linalg::cosine_similarity(x1, x2, 1, epsilon)
            .into_data()
            .assert_approx_eq::<FT>(&expected, tolerance);
    }
}