trustformers-core 0.1.1

Core traits and utilities for TrustformeRS
Documentation 
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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160

//! Tests for complex tensor functionality
//!
//! This module contains tests for the complex number support in the tensor system.

use crate::errors::Result;
use crate::tensor::Tensor;

#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_relative_eq;

    #[test]
    fn test_complex_tensor_creation() -> Result<()> {
        let real = vec![1.0, 2.0, 3.0, 4.0];
        let imag = vec![0.5, 1.5, 2.5, 3.5];
        let tensor = Tensor::complex(real, imag, &[2, 2])?;

        assert_eq!(tensor.shape(), vec![2, 2]);
        assert_eq!(tensor.dtype(), crate::tensor::DType::C32);
        Ok(())
    }

    #[test]
    fn test_complex_zeros() -> Result<()> {
        let tensor = Tensor::zeros_c32(&[3, 3])?;
        assert_eq!(tensor.shape(), vec![3, 3]);
        assert_eq!(tensor.dtype(), crate::tensor::DType::C32);
        Ok(())
    }

    #[test]
    fn test_real_and_imag_extraction() -> Result<()> {
        let real = vec![1.0, 2.0, 3.0, 4.0];
        let imag = vec![0.5, 1.5, 2.5, 3.5];
        let tensor = Tensor::complex(real.clone(), imag.clone(), &[2, 2])?;

        let real_part = tensor.real()?;
        let imag_part = tensor.imag()?;

        assert_eq!(real_part.data()?, real);
        assert_eq!(imag_part.data()?, imag);
        Ok(())
    }

    #[test]
    fn test_magnitude() -> Result<()> {
        let real = vec![3.0, 4.0];
        let imag = vec![4.0, 3.0];
        let tensor = Tensor::complex(real, imag, &[2])?;

        let magnitude = tensor.magnitude()?;
        let expected = [5.0, 5.0]; // sqrt(3^2 + 4^2) = sqrt(4^2 + 3^2) = 5

        for (actual, expected) in magnitude.data()?.iter().zip(expected.iter()) {
            assert_relative_eq!(actual, expected, epsilon = 1e-5);
        }
        Ok(())
    }

    #[test]
    fn test_phase() -> Result<()> {
        let real = vec![1.0, 0.0, -1.0, 0.0];
        let imag = vec![0.0, 1.0, 0.0, -1.0];
        let tensor = Tensor::complex(real, imag, &[4])?;

        let phase = tensor.phase()?;
        let expected = [
            0.0,
            std::f32::consts::FRAC_PI_2,
            std::f32::consts::PI,
            -std::f32::consts::FRAC_PI_2,
        ];

        for (actual, expected) in phase.data()?.iter().zip(expected.iter()) {
            assert_relative_eq!(actual, expected, epsilon = 1e-5);
        }
        Ok(())
    }

    #[test]
    fn test_complex_conjugate() -> Result<()> {
        let real = vec![1.0, 2.0];
        let imag = vec![3.0, 4.0];
        let tensor = Tensor::complex(real.clone(), imag.clone(), &[2])?;

        let conj = tensor.conj()?;
        let conj_real = conj.real()?;
        let conj_imag = conj.imag()?;

        assert_eq!(conj_real.data()?, real);
        assert_eq!(conj_imag.data()?, vec![-3.0, -4.0]);
        Ok(())
    }

    #[test]
    fn test_real_to_complex_conversion() -> Result<()> {
        let real_tensor = Tensor::from_vec(vec![1.0, 2.0, 3.0, 4.0], &[2, 2])?;
        let complex_tensor = real_tensor.to_complex()?;

        assert_eq!(complex_tensor.dtype(), crate::tensor::DType::C32);
        assert_eq!(complex_tensor.real()?.data()?, vec![1.0, 2.0, 3.0, 4.0]);
        assert_eq!(complex_tensor.imag()?.data()?, vec![0.0, 0.0, 0.0, 0.0]);
        Ok(())
    }

    #[test]
    fn test_complex_f64() -> Result<()> {
        let real: Vec<f64> = vec![1.0, 2.0, 3.0, 4.0];
        let imag: Vec<f64> = vec![0.5, 1.5, 2.5, 3.5];
        let tensor = Tensor::complex_f64(real.clone(), imag.clone(), &[2, 2])?;

        assert_eq!(tensor.shape(), vec![2, 2]);
        assert_eq!(tensor.dtype(), crate::tensor::DType::C64);

        let real_part = tensor.real()?;
        let imag_part = tensor.imag()?;

        // Note: conversion to f32 might have small differences
        for (actual, expected) in real_part.data()?.iter().zip(real.iter()) {
            assert_relative_eq!(*actual, *expected as f32, epsilon = 1e-5);
        }
        for (actual, expected) in imag_part.data()?.iter().zip(imag.iter()) {
            assert_relative_eq!(*actual, *expected as f32, epsilon = 1e-5);
        }
        Ok(())
    }

    #[test]
    fn test_real_tensor_magnitude() -> Result<()> {
        let tensor = Tensor::from_vec(vec![-3.0, 4.0, -5.0], &[3])?;
        let magnitude = tensor.magnitude()?;

        assert_eq!(magnitude.data()?, vec![3.0, 4.0, 5.0]);
        Ok(())
    }

    #[test]
    fn test_real_tensor_phase() -> Result<()> {
        let tensor = Tensor::from_vec(vec![1.0, -2.0, 3.0], &[3])?;
        let phase = tensor.phase()?;

        let expected = [0.0, std::f32::consts::PI, 0.0];
        let phase_data = phase.data()?;
        for (actual, expected) in phase_data.iter().zip(expected.iter()) {
            assert_relative_eq!(actual, expected, epsilon = 1e-5);
        }
        Ok(())
    }

    #[test]
    fn test_real_tensor_conjugate() -> Result<()> {
        let tensor = Tensor::from_vec(vec![1.0, 2.0, 3.0], &[3])?;
        let conj = tensor.conj()?;

        // For real tensors, conjugate should be the same
        assert_eq!(conj.data()?, tensor.data()?);
        Ok(())
    }
}