feat.tensor-complex

nodegen/helpers.py

@@ -10,6 +10,8 @@
 class FixedImpl(Enum):
     FP8x23 = 'FP8x23'
     FP16x16 = 'FP16x16'
+    FP64x64 = 'FP64x64'
+
 
 
 def to_fp(x: np.ndarray, fp_impl: FixedImpl):
@@ -18,15 +20,19 @@ def to_fp(x: np.ndarray, fp_impl: FixedImpl):
             return (x * 2**23).astype(np.int64)
         case FixedImpl.FP16x16:
             return (x * 2**16).astype(np.int64)
+        case FixedImpl.FP64x64:
+            return (x * 2**64)
 
 
 class Dtype(Enum):
     FP8x23 = 'FP8x23'
     FP16x16 = 'FP16x16'
+    FP64x64 = 'FP64x64'
     I8 = 'i8'
     I32 = 'i32'
     U32 = 'u32'
     BOOL = 'bool'
+    COMPLEX64 = 'complex64'
 
 
 class Tensor:
@@ -167,8 +173,15 @@ def get_data_statement(data: np.ndarray, dtype: Dtype) -> list[str]:
             return ["FP8x23 { "+f"mag: {abs(int(x))}, sign: {str(x < 0).lower()} "+"}" for x in data.flatten()]
         case Dtype.FP16x16:
             return ["FP16x16 { "+f"mag: {abs(int(x))}, sign: {str(x < 0).lower()} "+"}" for x in data.flatten()]
+        case Dtype.FP64x64:
+            return ["FP64x64 { "+f"mag: {abs(int(x))}, sign: {str(x < 0).lower()} "+"}" for x in data.flatten()]
         case Dtype.BOOL:
             return [str(x).lower() for x in data.flatten()]
+        case Dtype.COMPLEX64:
+            return ["complex64 { "+"real: FP64x64 { "+f"mag: {abs(int(np.real(x)))}, sign: {str(np.real(x) < 0).lower()} "+"} , img: FP64x64 { "+f"mag: {abs(int(np.imag(x)))}, sign: {str(np.imag(x) < 0).lower()} "+"} }" for x in data.flatten()]
+
+
+
 
 
 def get_data_statement_for_sequences(data: Sequence, dtype: Dtype) -> list[list[str]]:
@@ -238,6 +251,7 @@ def find_all_types(tensors: list[Tensor | Sequence]) -> list[Dtype]:
     Dtype.FP8x23: ["orion::operators::tensor::FP8x23Tensor",],
     Dtype.FP16x16: ["orion::operators::tensor::FP16x16Tensor",],
     Dtype.BOOL: ["orion::operators::tensor::BoolTensor",],
+    Dtype.COMPLEX64: ["orion::operators::tensor::Complex64Tensor",],
 }
 
 
@@ -266,6 +280,7 @@ def find_all_types(tensors: list[Tensor | Sequence]) -> list[Dtype]:
     Dtype.FP8x23: ["orion::operators::tensor::FP8x23TensorPartialEq",],
     Dtype.FP16x16: ["orion::operators::tensor::FP16x16TensorPartialEq",],
     Dtype.BOOL: ["orion::operators::tensor::BoolTensorPartialEq",],
+    Dtype.COMPLEX64: ["orion::operators::tensor::Complex64TensorPartialEq",],
 }
 
 
@@ -276,4 +291,5 @@ def find_all_types(tensors: list[Tensor | Sequence]) -> list[Dtype]:
     Dtype.FP8x23: ["orion::numbers::{FixedTrait, FP8x23}",],
     Dtype.FP16x16: ["orion::numbers::{FixedTrait, FP16x16}",],
     Dtype.BOOL: [],
+    Dtype.COMPLEX64: ["orion::numbers::{NumberTrait, complex64}",],
 }

nodegen/node/reduce_l2.py

@@ -4,6 +4,7 @@
 import numpy as np
 
 
+
 class Reduce_l2(RunAll):
     @staticmethod
     def reduce_l2_fp8x23():
@@ -107,4 +108,29 @@ def reduce_l2_axis_0():
 
         reduce_l2_export_do_not_keepdims()
         reduce_l2_export_keepdims()
+        reduce_l2_axis_0()
+
+    @staticmethod
+    def reduce_l2_complex64():
+
+
+
+        def reduce_l2_axis_0():
+            shape = [2, 3]
+            axes = np.array([0], dtype=np.int64)
+            keepdims = True
+            x = np.reshape(np.array([1.+2.j,  2.-1.j,  3.-3.j,  3.-2.j,  3.+5.j, 4.- 1.j]), shape)
+            y = np.sqrt(np.sum(a=np.square(abs(x)), axis=tuple(axes), keepdims=True))
+            print(to_fp(x.flatten(), FixedImpl.FP64x64))
+
+            x = Tensor(Dtype.COMPLEX64, x.shape, to_fp(
+            x.flatten(), FixedImpl.FP64x64))
+
+            y = Tensor(Dtype.COMPLEX64, y.shape, to_fp(
+            y.flatten(), FixedImpl.FP64x64))
+
+            name = "reduce_l2_complex64_axis_0"
+            make_test(
+                [x], y, "input_0.reduce_l2(0, true)", name)
+
         reduce_l2_axis_0()

src/numbers/complex_number/complex64.cairo

@@ -73,7 +73,12 @@ impl Complex64Impl of ComplexTrait<complex64, FP64x64> {
         let y = self.img;
         let two = FP64x64Impl::new(TWO, false);
         let real = (((x.pow(two) + y.pow(two)).sqrt() + x) / two).sqrt();
-        let img = (((x.pow(two) + y.pow(two)).sqrt() - x) / two).sqrt();
+        let img = if y == FP64x64Impl::ZERO() {
+            FP64x64Impl::ZERO()
+        } else {
+            (((x.pow(two) + y.pow(two)).sqrt() - x) / two).sqrt()
+        };
+
         let img = FP64x64Impl::new(img.mag, y.sign);
         complex64 { real, img }
     }

src/operators/tensor.cairo

@@ -39,3 +39,7 @@ use orion::operators::tensor::implementations::tensor_u32::{
 
 use orion::operators::tensor::implementations::tensor_bool::{BoolTensor, BoolTensorPartialEq};
 
+use orion::operators::tensor::implementations::tensor_complex64::{
+    Complex64Tensor, Complex64TensorAdd, Complex64TensorSub, Complex64TensorMul, Complex64TensorDiv,
+    Complex64TensorPartialEq,
+};

src/operators/tensor/implementations.cairo

@@ -8,3 +8,4 @@ mod tensor_fp64x64;
 mod tensor_fp32x32;
 mod tensor_fp16x16wide;
 mod tensor_fp8x23wide;
+mod tensor_complex64;