Merge pull request #599 from canacechan/BitShift

feat: bit shift

docs/framework/operators/tensor/README.md

@@ -133,6 +133,7 @@ use orion::operators::tensor::TensorTrait;
 | [`tensor.dynamic_quantize_linear`](tensor.dynamic\_quantize\_linear.md) | Computes the Scale, Zero Point and FP32->8Bit conversion of FP32 Input data.  |
 | [`tensor.scatter_nd`](tensor.scatter\_nd.md) | The output of the operation is produced by creating a copy of the input data, and then updating its value to values specified by updates at specific index positions specified by indices. Its output shape is the same as the shape of data |
 | [`tensor.label_encoder`](tensor.label\_encoder.md) | Maps each element in the input tensor to another value. |
+| [`tensor.bit_shift`](tensor.bit\_shift.md) | Bitwise shift operator performs element-wise operation. For each input element, if the attribute "direction" is "RIGHT", this operator moves its binary representation toward the right side so that the input value is effectively decreased. If the attribute "direction" is "LEFT", bits of binary representation moves toward the left side, which results the increase of its actual value.  |
 | [`tensor.eye_like`](tensor.eye\_like.md) | Generate a 2D tensor (matrix) with ones on the diagonal and zeros everywhere else. Only 2D tensors are supported, i.e. input T1 must be of rank 2. The shape of the output tensor is the same as the input tensor.  |
 
 ## Arithmetic Operations

docs/framework/operators/tensor/tensor.bit_shift.md

@@ -0,0 +1,68 @@
+# TensorTrait::bit_shift
+
+```rust
+        fn bit_shift(tensor1: @Tensor<T>, tensor2: @Tensor<T>, direction: felt252) -> Tensor<T>;
+```
+
+Bitwise shift operator performs element-wise operation. For each input element, if the attribute "direction" is "RIGHT", this operator moves its binary representation toward the right side so that the input value is effectively decreased. If the attribute "direction" is "LEFT", bits of binary representation moves toward the left side, which results the increase of its actual value. 
+The input tensor1 is the tensor to be shifted and another input tensor2 specifies the amounts of shifting. 
+
+## Args
+
+* `tensor1`(`@Tensor<T>`) - First operand, input to be shifted.
+* `tensor2`(`@Tensor<T>`) - Second operand, amounts of shift.
+* `direction`(@Tensor<T>) - Direction of moving bits. It can be either "RIGHT" (for right shift) or "LEFT" (for left shift).
+
+## Returns
+
+* Output tensor
+
+## Examples
+
+```rust
+
+use orion::operators::tensor::{U32Tensor, U32TensorAdd};
+use core::array::{ArrayTrait, SpanTrait};
+use orion::operators::tensor::{TensorTrait, Tensor};
+use orion::utils::{assert_eq, assert_seq_eq};
+use orion::operators::tensor::U32TensorPartialEq;
+
+
+fn example() -> Tensor<u32> {
+    let mut shape = ArrayTrait::<usize>::new();
+    shape.append(3);
+    shape.append(3);
+
+    let mut data = ArrayTrait::new();
+    data.append(21);
+    data.append(7);
+    data.append(18);
+    data.append(43);
+    data.append(49);
+    data.append(49);
+    data.append(4);
+    data.append(28);
+    data.append(24);
+    let input_0 = TensorTrait::new(shape.span(), data.span());
+
+    let mut shape1 = ArrayTrait::<usize>::new();
+    shape1.append(3);
+    shape1.append(3);
+
+    let mut data1 = ArrayTrait::new();
+    data1.append(4);
+    data1.append(0);
+    data1.append(0);
+    data1.append(3);
+    data1.append(1);
+    data1.append(1);
+    data1.append(1);
+    data1.append(1);
+    data1.append(0);
+    let input_1 = TensorTrait::new(shape1.span(), data1.span());
+
+
+    return TensorTrait::bit_shift(@input_0, @input_1, 'LEFT');
+}
+>>> [336 7 18 344 98 98 8 56 24]
+```

nodegen/node/bit_shift.py

@@ -0,0 +1,32 @@
+import numpy as np
+from nodegen.node import RunAll
+from ..helpers import make_test, to_fp, Tensor, Dtype, FixedImpl, Trait, get_data_statement
+
+class Bit_shift(RunAll):
+
+    @staticmethod
+    def left_u32():
+        x = np.random.randint(0, 60, (3, 3)).astype(np.uint32)
+        y = np.random.randint(0, 6, (3, 3)).astype(np.uint32)
+        z = x << y
+
+        x = Tensor(Dtype.U32, x.shape, x.flatten())
+        y = Tensor(Dtype.U32, y.shape, y.flatten())
+        z = Tensor(Dtype.U32, z.shape, z.flatten())
+
+        name = "bit_shift_left_u32"
+        make_test([x, y], z, "TensorTrait::bit_shift(@input_0, @input_1, 'LEFT')", name)
+
+    @staticmethod
+    def right_u32():
+        x = np.random.randint(0, 60, (3, 3)).astype(np.uint32)
+        y = np.random.randint(0, 6, (3, 3)).astype(np.uint32)
+        z = x >> y
+
+        x = Tensor(Dtype.U32, x.shape, x.flatten())
+        y = Tensor(Dtype.U32, y.shape, y.flatten())
+        z = Tensor(Dtype.U32, z.shape, z.flatten())
+
+        name = "bit_shift_right_u32"
+        make_test([x, y], z, "TensorTrait::bit_shift(@input_0, @input_1, 'RIGHT')", name)
+

src/operators/tensor/core.cairo

@@ -137,6 +137,7 @@ impl TensorSerde<T, impl TSerde: Serde<T>, impl TDrop: Drop<T>> of Serde<Tensor<
 /// scatter_nd - The output of the operation is produced by creating a copy of the input data, and then updating its value to values specified by updates at specific index positions specified by indices. Its output shape is the same as the shape of data
 /// center_crop_pad - Center crop or pad an input to given dimensions.
 /// label_encoder - Maps each element in the input tensor to another value.
+/// bit_shift - Bitwise shift operator performs element-wise operation. For each input element, if the attribute "direction" is "RIGHT", this operator moves its binary representation toward the right side so that the input value is effectively decreased. If the attribute "direction" is "LEFT", bits of binary representation moves toward the left side, which results the increase of its actual value. 
 /// eye_like - Generate a 2D tensor (matrix) with ones on the diagonal and zeros everywhere else. Only 2D tensors are supported, i.e. input T1 must be of rank 2. The shape of the output tensor is the same as the input tensor. 
 trait TensorTrait<T> {
     /// # tensor.new
@@ -6089,6 +6090,76 @@ trait TensorTrait<T> {
         values: Option<Span<T>>,
         values_tensor: Option<Tensor<T>>
     ) -> Tensor<T>;
+    /// # TensorTrait::bit_shift
+    ///
+    /// ```rust
+    ///         fn bit_shift(tensor1: @Tensor<T>, tensor2: @Tensor<T>, direction: felt252) -> Tensor<T>;
+    /// ```
+    ///
+    /// Bitwise shift operator performs element-wise operation. For each input element, if the attribute "direction" is "RIGHT", this operator moves its binary representation toward the right side so that the input value is effectively decreased. If the attribute "direction" is "LEFT", bits of binary representation moves toward the left side, which results the increase of its actual value. 
+    /// The input tensor1 is the tensor to be shifted and another input tensor2 specifies the amounts of shifting. 
+    ///
+    /// ## Args
+    ///
+    /// * `tensor1`(`@Tensor<T>`) - First operand, input to be shifted.
+    /// * `tensor2`(`@Tensor<T>`) - Second operand, amounts of shift.
+    /// * `direction`(@Tensor<T>) - Direction of moving bits. It can be either "RIGHT" (for right shift) or "LEFT" (for left shift).
+    ///
+    /// ## Returns
+    ///
+    /// * Output tensor
+    ///
+    /// ## Examples
+    ///
+    /// ```rust
+    /// 
+    /// use orion::operators::tensor::{U32Tensor, U32TensorAdd};
+    /// use core::array::{ArrayTrait, SpanTrait};
+    /// use orion::operators::tensor::{TensorTrait, Tensor};
+    /// use orion::utils::{assert_eq, assert_seq_eq};
+    /// use orion::operators::tensor::U32TensorPartialEq;
+    /// 
+    /// 
+    /// fn example() -> Tensor<u32> {
+    ///     let mut shape = ArrayTrait::<usize>::new();
+    ///     shape.append(3);
+    ///     shape.append(3);
+    /// 
+    ///     let mut data = ArrayTrait::new();
+    ///     data.append(21);
+    ///     data.append(7);
+    ///     data.append(18);
+    ///     data.append(43);
+    ///     data.append(49);
+    ///     data.append(49);
+    ///     data.append(4);
+    ///     data.append(28);
+    ///     data.append(24);
+    ///     let input_0 = TensorTrait::new(shape.span(), data.span());
+    /// 
+    ///     let mut shape1 = ArrayTrait::<usize>::new();
+    ///     shape1.append(3);
+    ///     shape1.append(3);
+    /// 
+    ///     let mut data1 = ArrayTrait::new();
+    ///     data1.append(4);
+    ///     data1.append(0);
+    ///     data1.append(0);
+    ///     data1.append(3);
+    ///     data1.append(1);
+    ///     data1.append(1);
+    ///     data1.append(1);
+    ///     data1.append(1);
+    ///     data1.append(0);
+    ///     let input_1 = TensorTrait::new(shape1.span(), data1.span());
+    /// 
+    /// 
+    ///     return TensorTrait::bit_shift(@input_0, @input_1, 'LEFT');
+    /// }
+    /// >>> [336 7 18 344 98 98 8 56 24]
+    /// ```
+    ///
+    fn bit_shift(tensor1: @Tensor<T>, tensor2: @Tensor<T>, direction: felt252) -> Tensor<T>;
     /// # TensorTrait::eye_like
     /// 
     /// ```rust

src/operators/tensor/implementations/tensor_bool.cairo

@@ -581,6 +581,12 @@ impl BoolTensor of TensorTrait<bool> {
         panic(array!['not supported!'])
     }
 
+    fn bit_shift(
+        tensor1: @Tensor<bool>, tensor2: @Tensor<bool>, direction: felt252
+    ) -> Tensor<bool> {
+        panic(array!['not supported!'])
+    }
+
     fn eye_like(self: @Tensor<bool>, k: Option<i32>) -> Tensor<bool> {
         panic(array!['not supported!'])
     }

src/operators/tensor/implementations/tensor_complex64.cairo

@@ -624,6 +624,12 @@ impl Complex64Tensor of TensorTrait<complex64> {
         panic(array!['not supported!'])
     }
 
+    fn bit_shift(
+        tensor1: @Tensor<complex64>, tensor2: @Tensor<complex64>, direction: felt252
+    ) -> Tensor<complex64> {
+        panic(array!['not supported!'])
+    }
+
     fn eye_like(self: @Tensor<complex64>, k: Option<i32>) -> Tensor<complex64> {
         panic(array!['not supported!'])
     }

src/operators/tensor/implementations/tensor_fp16x16.cairo

@@ -692,6 +692,12 @@ impl FP16x16Tensor of TensorTrait<FP16x16> {
         )
     }
 
+    fn bit_shift(
+        tensor1: @Tensor<FP16x16>, tensor2: @Tensor<FP16x16>, direction: felt252
+    ) -> Tensor<FP16x16> {
+        panic(array!['not supported!'])
+    }
+
     fn eye_like(self: @Tensor<FP16x16>, k: Option<i32>) -> Tensor<FP16x16> {
         math::eye_like::eye_like(self, k)
     }

src/operators/tensor/implementations/tensor_fp16x16wide.cairo

@@ -633,6 +633,12 @@ impl FP16x16WTensor of TensorTrait<FP16x16W> {
         )
     }
 
+    fn bit_shift(
+        tensor1: @Tensor<FP16x16W>, tensor2: @Tensor<FP16x16W>, direction: felt252
+    ) -> Tensor<FP16x16W> {
+        panic(array!['not supported!'])
+    }
+
     fn eye_like(self: @Tensor<FP16x16W>, k: Option<i32>) -> Tensor<FP16x16W> {
         math::eye_like::eye_like(self, k)
     }

src/operators/tensor/implementations/tensor_fp32x32.cairo

@@ -688,6 +688,12 @@ impl FP32x32Tensor of TensorTrait<FP32x32> {
         )
     }
 
+    fn bit_shift(
+        tensor1: @Tensor<FP32x32>, tensor2: @Tensor<FP32x32>, direction: felt252
+    ) -> Tensor<FP32x32> {
+        panic(array!['not supported!'])
+    }
+
     fn eye_like(self: @Tensor<FP32x32>, k: Option<i32>) -> Tensor<FP32x32> {
         math::eye_like::eye_like(self, k)
     }

src/operators/tensor/implementations/tensor_fp64x64.cairo

@@ -688,6 +688,12 @@ impl FP64x64Tensor of TensorTrait<FP64x64> {
         )
     }
 
+    fn bit_shift(
+        tensor1: @Tensor<FP64x64>, tensor2: @Tensor<FP64x64>, direction: felt252
+    ) -> Tensor<FP64x64> {
+        panic(array!['not supported!'])
+    }
+
     fn eye_like(self: @Tensor<FP64x64>, k: Option<i32>) -> Tensor<FP64x64> {
         math::eye_like::eye_like(self, k)
     }

src/operators/tensor/implementations/tensor_fp8x23.cairo

@@ -684,6 +684,12 @@ impl FP8x23Tensor of TensorTrait<FP8x23> {
         )
     }
 
+    fn bit_shift(
+        tensor1: @Tensor<FP8x23>, tensor2: @Tensor<FP8x23>, direction: felt252
+    ) -> Tensor<FP8x23> {
+        panic(array!['not supported!'])
+    }
+
     fn eye_like(self: @Tensor<FP8x23>, k: Option<i32>) -> Tensor<FP8x23> {
         math::eye_like::eye_like(self, k)
     }

src/operators/tensor/implementations/tensor_fp8x23wide.cairo

@@ -610,6 +610,12 @@ impl FP8x23WTensor of TensorTrait<FP8x23W> {
         )
     }
 
+    fn bit_shift(
+        tensor1: @Tensor<FP8x23W>, tensor2: @Tensor<FP8x23W>, direction: felt252
+    ) -> Tensor<FP8x23W> {
+        panic(array!['not supported!'])
+    }
+
     fn eye_like(self: @Tensor<FP8x23W>, k: Option<i32>) -> Tensor<FP8x23W> {
         math::eye_like::eye_like(self, k)
     }

src/operators/tensor/implementations/tensor_i32.cairo

@@ -652,6 +652,12 @@ impl I32Tensor of TensorTrait<i32> {
         )
     }
 
+    fn bit_shift(
+        tensor1: @Tensor<i32>, tensor2: @Tensor<i32>, direction: felt252
+    ) -> Tensor<i32> {
+        panic(array!['not supported!'])
+    }
+
     fn eye_like(self: @Tensor<i32>, k: Option<i32>) -> Tensor<i32> {
         math::eye_like::eye_like(self, k)
     }

src/operators/tensor/implementations/tensor_i8.cairo

@@ -655,6 +655,12 @@ impl I8Tensor of TensorTrait<i8> {
         )
     }
 
+    fn bit_shift(
+        tensor1: @Tensor<i8>, tensor2: @Tensor<i8>, direction: felt252
+    ) -> Tensor<i8> {
+        panic(array!['not supported!'])
+    }
+
     fn eye_like(self: @Tensor<i8>, k: Option<i32>) -> Tensor<i8> {
         math::eye_like::eye_like(self, k)
     }

src/operators/tensor/implementations/tensor_u32.cairo

@@ -580,6 +580,12 @@ impl U32Tensor of TensorTrait<u32> {
         )
     }
 
+    fn bit_shift(
+        tensor1: @Tensor<u32>, tensor2: @Tensor<u32>, direction: felt252
+    ) -> Tensor<u32> {
+        math::bit_shift::bit_shift(tensor1, tensor2, direction)
+    }
+
     fn eye_like(self: @Tensor<u32>, k: Option<i32>) -> Tensor<u32> {
         math::eye_like::eye_like(self, k)
     }

src/operators/tensor/math.cairo

@@ -68,4 +68,5 @@ mod hann_window;
 mod hamming_window;
 mod blackman_window;
 mod scatter_nd;
+mod bit_shift;
 mod eye_like;