Merge pull request #485 from gizatechxyz/refactor-tensor-and

Refactoring: Tensor AND return BoolTensor

docs/framework/operators/tensor/README.md

@@ -65,7 +65,6 @@ use orion::operators::tensor::TensorTrait;
 | [`tensor.log`](tensor.log.md) | Computes the natural log of all elements of the input tensor. |
 | [`tensor.abs`](tensor.abs.md) | Computes the absolute value of all elements in the input tensor. |
 | [`tensor.neg`](tensor.neg.md) | Computes the negation of all elements in the input tensor. |
-| [`tensor.not`](tensor.not.md) | Computes the logical negation of all elements in the input tensor. |
 | [`tensor.ceil`](tensor.ceil.md) | Rounds up the value of each element in the input tensor. |
 | [`tensor.sqrt`](tensor.sqrt.md) | Computes the square root of all elements of the input tensor. |
 | [`tensor.sin`](tensor.sin.md) | Computes the sine of all elements of the input tensor. |
@@ -122,6 +121,7 @@ use orion::operators::tensor::TensorTrait;
 | [`tensor.concat_from_sequence`](tensor.concat\_from\_sequence.md) | Concatenate a sequence of tensors into a single tensor. |
 | [`tensor.is_nan`](tensor.is\_nan.md) | Returns which elements of the input are NaN. |
 | [`tensor.is_inf`](tensor.is\_inf.md) | Maps infinity to true and other values to false. |
+| [`tensor.not`](tensor.not.md) | Computes the logical negation of all elements in the input tensor. |
 
 ## Arithmetic Operations
 

nodegen/node/and.py

@@ -13,7 +13,7 @@ def default():
 
             x = Tensor(Dtype.U32, x.shape, x.flatten())
             y = Tensor(Dtype.U32, y.shape, y.flatten())
-            z = Tensor(Dtype.U32, z.shape, z.flatten())
+            z = Tensor(Dtype.BOOL, z.shape, z.flatten())
 
             name = "and_u32"
             make_test([x, y], z, "input_0.and(@input_1)", name)
@@ -25,7 +25,7 @@ def broadcast():
 
             x = Tensor(Dtype.U32, x.shape, x.flatten())
             y = Tensor(Dtype.U32, y.shape, y.flatten())
-            z = Tensor(Dtype.U32, z.shape, z.flatten())
+            z = Tensor(Dtype.BOOL, z.shape, z.flatten())
 
             name = "and_u32_broadcast"
             make_test([x, y], z, "input_0.and(@input_1)", name)
@@ -42,7 +42,7 @@ def default():
 
             x = Tensor(Dtype.I32, x.shape, x.flatten())
             y = Tensor(Dtype.I32, y.shape, y.flatten())
-            z = Tensor(Dtype.U32, z.shape, z.flatten())
+            z = Tensor(Dtype.BOOL, z.shape, z.flatten())
 
             name = "and_i32"
             make_test([x, y], z, "input_0.and(@input_1)", name)
@@ -54,7 +54,7 @@ def broadcast():
 
             x = Tensor(Dtype.I32, x.shape, x.flatten())
             y = Tensor(Dtype.I32, y.shape, y.flatten())
-            z = Tensor(Dtype.U32, z.shape, z.flatten())
+            z = Tensor(Dtype.BOOL, z.shape, z.flatten())
 
             name = "and_i32_broadcast"
             make_test([x, y], z, "input_0.and(@input_1)", name)
@@ -71,7 +71,7 @@ def default():
 
             x = Tensor(Dtype.I8, x.shape, x.flatten())
             y = Tensor(Dtype.I8, y.shape, y.flatten())
-            z = Tensor(Dtype.U32, z.shape, z.flatten())
+            z = Tensor(Dtype.BOOL, z.shape, z.flatten())
 
             name = "and_i8"
             make_test([x, y], z, "input_0.and(@input_1)", name)
@@ -83,7 +83,7 @@ def broadcast():
 
             x = Tensor(Dtype.I8, x.shape, x.flatten())
             y = Tensor(Dtype.I8, y.shape, y.flatten())
-            z = Tensor(Dtype.U32, z.shape, z.flatten())
+            z = Tensor(Dtype.BOOL, z.shape, z.flatten())
 
             name = "and_i8_broadcast"
             make_test([x, y], z, "input_0.and(@input_1)", name)
@@ -102,7 +102,7 @@ def default():
                 x.flatten(), FixedImpl.FP8x23))
             y = Tensor(Dtype.FP8x23, y.shape,  to_fp(
                 y.flatten(), FixedImpl.FP8x23))
-            z = Tensor(Dtype.U32, z.shape, z.flatten())
+            z = Tensor(Dtype.BOOL, z.shape, z.flatten())
 
             name = "and_fp8x23"
             make_test([x, y], z, "input_0.and(@input_1)", name)
@@ -116,7 +116,7 @@ def broadcast():
                 x.flatten(), FixedImpl.FP8x23))
             y = Tensor(Dtype.FP8x23, y.shape,  to_fp(
                 y.flatten(), FixedImpl.FP8x23))
-            z = Tensor(Dtype.U32, z.shape, z.flatten())
+            z = Tensor(Dtype.BOOL, z.shape, z.flatten())
 
             name = "and_fp8x23_broadcast"
             make_test([x, y], z, "input_0.and(@input_1)", name)
@@ -135,7 +135,7 @@ def default():
                 x.flatten(), FixedImpl.FP16x16))
             y = Tensor(Dtype.FP16x16, y.shape,  to_fp(
                 y.flatten(), FixedImpl.FP16x16))
-            z = Tensor(Dtype.U32, z.shape, z.flatten())
+            z = Tensor(Dtype.BOOL, z.shape, z.flatten())
 
             name = "and_fp16x16"
             make_test([x, y], z, "input_0.and(@input_1)", name)
@@ -149,7 +149,7 @@ def broadcast():
                 x.flatten(), FixedImpl.FP16x16))
             y = Tensor(Dtype.FP16x16, y.shape,  to_fp(
                 y.flatten(), FixedImpl.FP16x16))
-            z = Tensor(Dtype.U32, z.shape, z.flatten())
+            z = Tensor(Dtype.BOOL, z.shape, z.flatten())
 
             name = "and_fp16x16_broadcast"
             make_test([x, y], z, "input_0.and(@input_1)", name)

src/operators/tensor/core.cairo

@@ -3452,7 +3452,7 @@ trait TensorTrait<T> {
     /// #tensor.and
     ///
     /// ```rust
-    ///     fn and(self: @Tensor<T>, other: @Tensor<T>) -> Tensor<usize>;
+    ///     fn and(self: @Tensor<T>, other: @Tensor<T>) -> Tensor<bool>;
     /// ```
     ///
     /// Computes the logical AND of two tensors element-wise.
@@ -3471,7 +3471,7 @@ trait TensorTrait<T> {
     ///
     /// ## Returns
     ///
-    /// A new `Tensor<usize>` of booleans (0 or 1) with the same shape as the broadcasted inputs.
+    /// A new `Tensor<bool>` with the same shape as the broadcasted inputs.
     ///
     /// ## Examples
     ///
@@ -3482,7 +3482,7 @@ trait TensorTrait<T> {
     /// 
     /// use orion::operators::tensor::{TensorTrait, Tensor, U32Tensor};
     /// 
-    /// fn and_example() -> Tensor<usize> {
+    /// fn and_example() -> Tensor<bool> {
     ///     let tensor_1 = TensorTrait::<u32>::new(
     ///         shape: array![3, 3].span(), data: array![0, 1, 2, 3, 4, 5, 6, 7, 8].span(),
     ///     );
@@ -3493,7 +3493,7 @@ trait TensorTrait<T> {
     /// 
     ///     return tensor_1.and(@tensor_2);
     /// }
-    /// >>> [0,1,1,0,1,1,0,1,1]
+    /// >>> [false, true, true, false, true, true, false, true, true]
     /// ```
     ///
     /// Case 2: Compare tensors with different shapes
@@ -3503,7 +3503,7 @@ trait TensorTrait<T> {
     /// 
     /// use orion::operators::tensor::{TensorTrait, Tensor, U32Tensor};
     /// 
-    /// fn and_example() -> Tensor<usize> {
+    /// fn and_example() -> Tensor<bool> {
     ///     let tensor_1 = TensorTrait::<u32>::new(
     ///         shape: array![3, 3].span(), data: array![0, 1, 2, 3, 4, 5, 6, 7, 8].span(),
     ///     );
@@ -3514,10 +3514,10 @@ trait TensorTrait<T> {
     /// 
     ///     return tensor_1.and(@tensor_2);
     /// }
-    /// >>> [0,1,1,0,1,1,0,1,1]
+    /// >>> [false, true, true, false, true, true, false, true, true]
     /// ```
     ///
-    fn and(self: @Tensor<T>, other: @Tensor<T>) -> Tensor<usize>;
+    fn and(self: @Tensor<T>, other: @Tensor<T>) -> Tensor<bool>;
     /// #tensor.where
     ///
     /// ```rust
@@ -4730,7 +4730,9 @@ trait TensorTrait<T> {
     /// >>> [false, false, true, false, true, true]
     /// ```
     ///
-    fn is_inf(self: @Tensor<T>, detect_negative: Option<u8>, detect_positive: Option<u8>) -> Tensor<bool>;
+    fn is_inf(
+        self: @Tensor<T>, detect_negative: Option<u8>, detect_positive: Option<u8>
+    ) -> Tensor<bool>;
     /// ## tensor.is_nan
     ///
     /// ```rust
@@ -4830,7 +4832,9 @@ trait TensorTrait<T> {
     /// >>> (2, 4 ) 
     /// ```
     ///
-    fn concat_from_sequence(sequence: Array<Tensor<T>>, axis: i32, new_axis: Option<usize>) -> Tensor<T>;
+    fn concat_from_sequence(
+        sequence: Array<Tensor<T>>, axis: i32, new_axis: Option<usize>
+    ) -> Tensor<T>;
     /// #tensor.not
     /// 
     /// ```rust
@@ -5433,4 +5437,4 @@ fn clip<
 /// Cf: TensorTrait::identity docstring
 fn identity<T>(self: @Tensor<T>) -> Tensor<T> {
     Tensor::<T> { shape: *self.shape, data: *self.data }
-}
+}

src/operators/tensor/implementations/tensor_bool.cairo

@@ -256,7 +256,7 @@ impl BoolTensor of TensorTrait<bool> {
         panic(array!['not supported!'])
     }
 
-    fn and(self: @Tensor<bool>, other: @Tensor<bool>) -> Tensor<usize> {
+    fn and(self: @Tensor<bool>, other: @Tensor<bool>) -> Tensor<bool> {
         panic(array!['not supported!'])
     }
 

src/operators/tensor/implementations/tensor_fp16x16.cairo

@@ -10,7 +10,7 @@ use orion::operators::tensor::core::{
 };
 use orion::operators::tensor::{math, linalg, quantization, core, ml};
 use orion::numbers::{i8, i32, NumberTrait, FP16x16};
-use orion::operators::tensor::implementations::{tensor_i8::I8Tensor, tensor_u32::U32Tensor};
+use orion::operators::tensor::implementations::{tensor_i8::I8Tensor, tensor_u32::U32Tensor, tensor_bool::BoolTensor};
 
 impl FP16x16Tensor of TensorTrait<FP16x16> {
     fn new(shape: Span<usize>, data: Span<FP16x16>) -> Tensor<FP16x16> {
@@ -372,7 +372,7 @@ impl FP16x16Tensor of TensorTrait<FP16x16> {
         core::clip(self, min, max)
     }
 
-    fn and(self: @Tensor<FP16x16>, other: @Tensor<FP16x16>) -> Tensor<usize> {
+    fn and(self: @Tensor<FP16x16>, other: @Tensor<FP16x16>) -> Tensor<bool> {
         math::and::and(self, other)
     }
 

src/operators/tensor/implementations/tensor_fp16x16wide.cairo

@@ -10,7 +10,7 @@ use orion::operators::tensor::core::{
 };
 use orion::operators::tensor::{math, linalg, quantization, core, ml};
 use orion::numbers::{i8, i32, NumberTrait, FP16x16W};
-use orion::operators::tensor::implementations::{tensor_i8::I8Tensor, tensor_u32::U32Tensor};
+use orion::operators::tensor::implementations::{tensor_i8::I8Tensor, tensor_u32::U32Tensor, tensor_bool::BoolTensor};
 
 impl FP16x16WTensor of TensorTrait<FP16x16W> {
     fn new(shape: Span<usize>, data: Span<FP16x16W>) -> Tensor<FP16x16W> {
@@ -334,7 +334,7 @@ impl FP16x16WTensor of TensorTrait<FP16x16W> {
         core::clip(self, min, max)
     }
 
-    fn and(self: @Tensor<FP16x16W>, other: @Tensor<FP16x16W>) -> Tensor<usize> {
+    fn and(self: @Tensor<FP16x16W>, other: @Tensor<FP16x16W>) -> Tensor<bool> {
         math::and::and(self, other)
     }
 

src/operators/tensor/implementations/tensor_fp32x32.cairo

@@ -11,7 +11,7 @@ use orion::operators::tensor::core::{
 use orion::operators::tensor::{math, linalg, quantization, core, ml};
 use orion::numbers::{i8, i32, NumberTrait, FP32x32, FP32x32Impl};
 use orion::numbers::fixed_point::implementations::fp32x32::core::ONE;
-use orion::operators::tensor::implementations::{tensor_i8::I8Tensor, tensor_u32::U32Tensor};
+use orion::operators::tensor::implementations::{tensor_i8::I8Tensor, tensor_u32::U32Tensor, tensor_bool::BoolTensor};
 
 impl FP32x32Tensor of TensorTrait<FP32x32> {
     fn new(shape: Span<usize>, data: Span<FP32x32>) -> Tensor<FP32x32> {
@@ -373,7 +373,7 @@ impl FP32x32Tensor of TensorTrait<FP32x32> {
         core::clip(self, min, max)
     }
 
-    fn and(self: @Tensor<FP32x32>, other: @Tensor<FP32x32>) -> Tensor<usize> {
+    fn and(self: @Tensor<FP32x32>, other: @Tensor<FP32x32>) -> Tensor<bool> {
         math::and::and(self, other)
     }
 

src/operators/tensor/implementations/tensor_fp64x64.cairo

@@ -11,7 +11,7 @@ use orion::operators::tensor::core::{
 use orion::operators::tensor::{math, linalg, quantization, core, ml};
 use orion::numbers::{i8, i32, NumberTrait, FP64x64, FP64x64Impl};
 use orion::numbers::fixed_point::implementations::fp64x64::core::ONE;
-use orion::operators::tensor::implementations::{tensor_i8::I8Tensor, tensor_u32::U32Tensor};
+use orion::operators::tensor::implementations::{tensor_i8::I8Tensor, tensor_u32::U32Tensor, tensor_bool::BoolTensor};
 
 impl FP64x64Tensor of TensorTrait<FP64x64> {
     fn new(shape: Span<usize>, data: Span<FP64x64>) -> Tensor<FP64x64> {
@@ -373,7 +373,7 @@ impl FP64x64Tensor of TensorTrait<FP64x64> {
         core::clip(self, min, max)
     }
 
-    fn and(self: @Tensor<FP64x64>, other: @Tensor<FP64x64>) -> Tensor<usize> {
+    fn and(self: @Tensor<FP64x64>, other: @Tensor<FP64x64>) -> Tensor<bool> {
         math::and::and(self, other)
     }
 

src/operators/tensor/implementations/tensor_fp8x23.cairo

@@ -10,7 +10,7 @@ use orion::operators::tensor::core::{
 };
 use orion::operators::tensor::{math, linalg, quantization, core, ml};
 use orion::numbers::{i8, i32, NumberTrait, FP8x23};
-use orion::operators::tensor::implementations::{tensor_i8::I8Tensor, tensor_u32::U32Tensor};
+use orion::operators::tensor::implementations::{tensor_i8::I8Tensor, tensor_u32::U32Tensor, tensor_bool::BoolTensor};
 
 impl FP8x23Tensor of TensorTrait<FP8x23> {
     fn new(shape: Span<usize>, data: Span<FP8x23>) -> Tensor<FP8x23> {
@@ -373,7 +373,7 @@ impl FP8x23Tensor of TensorTrait<FP8x23> {
         core::clip(self, min, max)
     }
 
-    fn and(self: @Tensor<FP8x23>, other: @Tensor<FP8x23>) -> Tensor<usize> {
+    fn and(self: @Tensor<FP8x23>, other: @Tensor<FP8x23>) -> Tensor<bool> {
         math::and::and(self, other)
     }
 

src/operators/tensor/implementations/tensor_fp8x23wide.cairo

@@ -10,7 +10,7 @@ use orion::operators::tensor::core::{
 };
 use orion::operators::tensor::{math, linalg, quantization, core, ml};
 use orion::numbers::{i8, i32, NumberTrait, FP8x23W};
-use orion::operators::tensor::implementations::{tensor_i8::I8Tensor, tensor_u32::U32Tensor};
+use orion::operators::tensor::implementations::{tensor_i8::I8Tensor, tensor_u32::U32Tensor, tensor_bool::BoolTensor};
 
 impl FP8x23WTensor of TensorTrait<FP8x23W> {
     fn new(shape: Span<usize>, data: Span<FP8x23W>) -> Tensor<FP8x23W> {
@@ -323,7 +323,7 @@ impl FP8x23WTensor of TensorTrait<FP8x23W> {
         core::clip(self, min, max)
     }
 
-    fn and(self: @Tensor<FP8x23W>, other: @Tensor<FP8x23W>) -> Tensor<usize> {
+    fn and(self: @Tensor<FP8x23W>, other: @Tensor<FP8x23W>) -> Tensor<bool> {
         math::and::and(self, other)
     }
 

src/operators/tensor/implementations/tensor_i32.cairo

@@ -10,7 +10,7 @@ use orion::operators::tensor::core::{
 };
 use orion::operators::tensor::{math, linalg, quantization, core, ml};
 use orion::numbers::{i32, i8, NumberTrait};
-use orion::operators::tensor::implementations::{tensor_u32::U32Tensor, tensor_i8::I8Tensor};
+use orion::operators::tensor::implementations::{tensor_u32::U32Tensor, tensor_i8::I8Tensor, tensor_bool::BoolTensor};
 
 
 impl I32Tensor of TensorTrait<i32> {
@@ -371,7 +371,7 @@ impl I32Tensor of TensorTrait<i32> {
         core::clip(self, min, max)
     }
 
-    fn and(self: @Tensor<i32>, other: @Tensor<i32>) -> Tensor<usize> {
+    fn and(self: @Tensor<i32>, other: @Tensor<i32>) -> Tensor<bool> {
         math::and::and(self, other)
     }
 

src/operators/tensor/implementations/tensor_i8.cairo

@@ -10,7 +10,7 @@ use orion::operators::tensor::core::{
 };
 use orion::operators::tensor::{math, linalg, quantization, core, ml};
 use orion::numbers::{i8, i32, NumberTrait};
-use orion::operators::tensor::implementations::tensor_u32::U32Tensor;
+use orion::operators::tensor::implementations::{tensor_u32::U32Tensor, tensor_bool::BoolTensor};
 
 impl I8Tensor of TensorTrait<i8> {
     fn new(shape: Span<usize>, data: Span<i8>) -> Tensor<i8> {
@@ -371,7 +371,7 @@ impl I8Tensor of TensorTrait<i8> {
         core::clip(self, min, max)
     }
 
-    fn and(self: @Tensor<i8>, other: @Tensor<i8>) -> Tensor<usize> {
+    fn and(self: @Tensor<i8>, other: @Tensor<i8>) -> Tensor<bool> {
         math::and::and(self, other)
     }
 

src/operators/tensor/implementations/tensor_u32.cairo

@@ -10,7 +10,7 @@ use orion::operators::tensor::core::{
 };
 use orion::operators::tensor::{math, linalg, quantization, core, ml};
 use orion::numbers::{i8, i32, NumberTrait};
-use orion::operators::tensor::implementations::tensor_i8::I8Tensor;
+use orion::operators::tensor::implementations::{tensor_i8::I8Tensor, tensor_bool::BoolTensor};
 
 impl U32Tensor of TensorTrait<u32> {
     fn new(shape: Span<usize>, data: Span<u32>) -> Tensor<u32> {
@@ -315,7 +315,7 @@ impl U32Tensor of TensorTrait<u32> {
         core::clip(self, min, max)
     }
 
-    fn and(self: @Tensor<u32>, other: @Tensor<u32>) -> Tensor<usize> {
+    fn and(self: @Tensor<u32>, other: @Tensor<u32>) -> Tensor<bool> {
         math::and::and(self, other)
     }
 

src/operators/tensor/math/and.cairo

@@ -13,14 +13,14 @@ fn and<
     T,
     MAG,
     impl TNumber: NumberTrait<T, MAG>,
-    impl UsizeFTensor: TensorTrait<usize>,
+    impl UsizeFTensor: TensorTrait<bool>,
     impl TCopy: Copy<T>,
     impl TDrop: Drop<T>
 >(
     y: @Tensor<T>, z: @Tensor<T>
-) -> Tensor<usize> {
+) -> Tensor<bool> {
     let broadcasted_shape = broadcast_shape(*y.shape, *z.shape);
-    let mut result: Array<usize> = ArrayTrait::new();
+    let mut result: Array<bool> = ArrayTrait::new();
 
     let num_elements = len_from_shape(broadcasted_shape);
 
@@ -32,9 +32,9 @@ fn and<
         let indices_other = broadcast_index_mapping(*z.shape, indices_broadcasted);
 
         if NumberTrait::and(*(*y.data)[indices_self], *(*z.data)[indices_other]) {
-            result.append(1);
+            result.append(true);
         } else {
-            result.append(0);
+            result.append(false);
         }
 
         n += 1;