Merge pull request #621 from gizatechxyz/refactoring

Refactoring

docs/framework/operators/neural-network/nn.softmax.md

@@ -1,7 +1,7 @@
 # NNTrait::softmax
 
 ```rust 
-   fn softmax(tensor: @Tensor<T>, axis: usize) -> Tensor<T>;
+   fn softmax(tensor: @Tensor<T>, axis: Option<i32>) -> Tensor<T>;
 ```
 
 Applies the Softmax function to an n-dimensional input Tensor rescaling them so that the elements of the n-dimensional output Tensor lie in the range \[0,1] and sum to 1.
@@ -13,7 +13,7 @@ $$
 ## Args
 
 * `tensor`(`@Tensor<T>`) - The input tensor.
-* `axis`(`usize`) - The axis along which to compute the softmax.
+* `axis`(`Option<i32>`) - Describes the dimension Softmax will be performed on. Negative value means counting dimensions from the back. Accepted range is [-r, r-1] where r = rank(input).
 
 ## Returns
 
@@ -44,7 +44,7 @@ fn softmax_example() -> Tensor<FP8x23> {
             .span(),
     );
 
-    return NNTrait::softmax(@tensor, 1);
+    return NNTrait::softmax(@tensor, Option::Some(1));
 }
 >>> [[2255697,6132911],[2255697,6132911]]
     // The fixed point representation of

docs/framework/operators/tensor/tensor.argmax.md

@@ -1,15 +1,15 @@
 # tensor.argmax
 
 ```rust 
-   fn argmax(self: @Tensor<T>, axis: usize, keepdims: Option<bool>, select_last_index: Option<bool>) -> Tensor<usize>;
+   fn argmax(self: @Tensor<T>, axis: i32, keepdims: Option<bool>, select_last_index: Option<bool>) -> Tensor<i32>;
 ```
 
 Returns the index of the maximum value along the specified axis.
 
 ## Args
 
 * `self`(`@Tensor<T>`) - The input tensor.
-* `axis`(`usize`) - The axis along which to compute the argmax.
+* `axis`(`i32`) - The axis along which to compute the argmax.
 * `keepdims`(`Option<bool>`) - If true, retains reduced dimensions with length 1. Defaults to true.
 * `select_last_index`(`Option<bool>`) - If true, the index of the last occurrence of the maximum value is returned. Defaults to false.   
 

docs/framework/operators/tensor/tensor.gather.md

@@ -1,16 +1,16 @@
 # tensor.gather
 
 ```rust 
-   fn gather(self: @Tensor<T>, indices: Tensor<T>, axis: Option<usize>) -> Tensor<T>;
+   fn gather(self: @Tensor<T>, indices: Tensor<i32>, axis: Option<i32>) -> Tensor<T>;
 ```
 
 Gather entries of the axis dimension of data.
 
 ## Args
 
 * `self`(`@Tensor<T>`) - The input tensor.
-* `indices`(`Tensor<T>`) - Tensor of indices.
-* `axis`(`Option<usize>`) - Axis to gather on. Default: axis=0.
+* `indices`(`Tensor<i32>`) - Tensor of indices.
+* `axis`(`Option<i32>`) - Axis to gather on. Default: axis=0.
 
 ## Panics
 
@@ -32,7 +32,7 @@ fn gather_example() -> Tensor<u32> {
         shape: array![2, 3].span(), 
         data: array![[ 1, 2, 3],[4, 5, 6]].span(), 
     );
-    let indices = TensorTrait::<u32>::new(
+    let indices = TensorTrait::<i32>::new(
         shape: array![1, 1].span(), 
         data: array![1, 0].span(), 
     );

docs/framework/operators/tensor/tensor.gather_elements.md

@@ -1,16 +1,16 @@
 # tensor.gather_elements
 
 ```rust 
-   fn gather_elements(self: @Tensor<T>, indices: Tensor<T>, axis: Option<usize>) -> Tensor<T>;
+   fn gather_elements(self: @Tensor<T>, indices: Tensor<i32>, axis: Option<i32>) -> Tensor<T>;
 ```
 
 GatherElements is an indexing operation that produces its output by indexing into the input data tensor at index positions determined by elements of the indices tensor.
 
 ## Args
 
 * `self`(`@Tensor<T>`) - The input tensor.
-* `indices`(`Tensor<T>`) - Tensor of indices.
-* `axis`(`Option<usize>`) - Axis to gather_elements on. Default: axis=0.
+* `indices`(`Tensor<i32>`) - Tensor of indices.
+* `axis`(`Option<i32>`) - Axis to gather_elements on. Default: axis=0.
 
 ## Panics
 
@@ -32,7 +32,7 @@ fn gather_elements_example() -> Tensor<u32> {
         shape: array![3, 3].span(), 
         data: array![[ 1, 2, 3],[4, 5, 6], [7, 8, 9]].span(), 
     );
-    let indices = TensorTrait::<u32>::new(
+    let indices = TensorTrait::<i32>::new(
         shape: array![1, 2, 0].span(), 
         data: array![2, 0, 0].span(), 
     );

docs/framework/operators/tensor/tensor.less.md

@@ -1,7 +1,7 @@
 #tensor.less
 
 ```rust
-    fn less(self: @Tensor<T>, other: @Tensor<T>) -> Tensor<usize>;
+    fn less(self: @Tensor<T>, other: @Tensor<T>) -> Tensor<i32>;
 ```
 
 Check if each element of the first tensor is less than the corresponding element of the second tensor.
@@ -20,7 +20,7 @@ The input tensors must have either:
 
 ## Returns
 
-A new `Tensor<usize>` of booleans (0 or 1) with the same shape as the broadcasted inputs.
+A new `Tensor<bool>` of booleans with the same shape as the broadcasted inputs.
 
 ## Examples
 
@@ -31,7 +31,7 @@ use core::array::{ArrayTrait, SpanTrait};
 
 use orion::operators::tensor::{TensorTrait, Tensor, U32Tensor};
 
-fn less_example() -> Tensor<usize> {
+fn less_example() -> Tensor<i32> {
     let tensor_1 = TensorTrait::<u32>::new(
         shape: array![3, 3, 3].span(), data: array![0, 1, 2, 3, 4, 5, 6, 7, 8].span(),
     );
@@ -53,7 +53,7 @@ use core::array::{ArrayTrait, SpanTrait};
 
 use orion::operators::tensor::{TensorTrait, Tensor, U32Tensor};
 
-fn less_example() -> Tensor<usize> {
+fn less_example() -> Tensor<i32> {
     let tensor_1 = TensorTrait::<u32>::new(
         shape: array![3, 3, 3].span(), data: array![0, 1, 2, 3, 4, 5, 6, 7, 8].span(),
     );

docs/framework/operators/tensor/tensor.less_equal.md

@@ -1,7 +1,7 @@
 #tensor.less_equal
 
 ```rust
-    fn less_equal(self: @Tensor<T>, other: @Tensor<T>) -> Tensor<usize>;
+    fn less_equal(self: @Tensor<T>, other: @Tensor<T>) -> Tensor<i32>;
 ```
 
 Check if each element of the first tensor is less than or equal to the corresponding element of the second tensor.
@@ -20,7 +20,7 @@ The input tensors must have either:
 
 ## Returns
 
-A new `Tensor<usize>` of booleans (0 or 1) with the same shape as the broadcasted inputs.
+A new `Tensor<i32>` of booleans (0 or 1) with the same shape as the broadcasted inputs.
 
 ## Examples
 
@@ -31,7 +31,7 @@ use core::array::{ArrayTrait, SpanTrait};
 
 use orion::operators::tensor::{TensorTrait, Tensor, U32Tensor};
 
-fn less_equal_example() -> Tensor<usize> {
+fn less_equal_example() -> Tensor<i32> {
     let tensor_1 = TensorTrait::<u32>::new(
         shape: array![3, 3, 3].span(), data: array![0, 1, 2, 3, 4, 5, 6, 7, 8].span(),
     );
@@ -53,7 +53,7 @@ use core::array::{ArrayTrait, SpanTrait};
 
 use orion::operators::tensor::{TensorTrait, Tensor, U32Tensor};
 
-fn less_equal_example() -> Tensor<usize> {
+fn less_equal_example() -> Tensor<i32> {
     let tensor_1 = TensorTrait::<u32>::new(
         shape: array![3, 3, 3].span(), data: array![0, 1, 2, 3, 4, 5, 6, 7, 8].span(),
     );

docs/framework/operators/tensor/tensor.reduce_sum.md

@@ -1,24 +1,21 @@
 ## tensor.reduce_sum
 
 ```rust 
-   fn reduce_sum(self: @Tensor<T>, axis: usize, keepdims: bool) -> Tensor<T>;
+   fn reduce_sum(self: @Tensor<T>, axes: Option<Span<i32>>, keepdims: Option<bool>, noop_with_empty_axes: Option<bool>) -> Tensor<T>;
 ```
 
 Reduces a tensor by summing its elements along a specified axis.
 
 ## Args
 
 * `self`(`@Tensor<T>`) - The input tensor.
-* `axis`(`usize`) - The dimension to reduce.
-* `keepdims`(`bool`) - If true, retains reduced dimensions with length 1.
-
-## Panics 
-
-* Panics if axis is not in the range of the input tensor's dimensions.
+* `axes`(`Option<Span<i32>>`) - Optional input list of integers, along which to reduce. The default is to reduce over all the dimensions of the input tensor if 'noop_with_empty_axes' is false, else act as an Identity op when 'noop_with_empty_axes' is true.
+* `keepdims`(`Option<bool>`) - Keep the reduced dimension or not, default 1 means keep reduced dimension.
+* `noop_with_empty_axes`(`Option<bool>`) - Defines behavior if 'axes' is empty. Default behavior with 'false' is to reduce all axes. When axes is empty and this attribute is set to true, input tensor will not be reduced,and the output tensor would be equivalent to input tensor.
 
 ## Returns
 
-A new `Tensor<T>` instance with the specified axis reduced by summing its elements.
+Reduced output tensor.
 
 ## Examples
 
@@ -33,7 +30,7 @@ fn reduce_sum_example() -> Tensor<u32> {
     );
 
     // We can call `reduce_sum` function as follows.
-    return tensor.reduce_sum(axis: 0, keepdims: false);
+    return tensor.reduce_sum(axes: Option::None, keepdims: false);
 }
 >>> [[4,6],[8,10]]
 ```

docs/framework/operators/tensor/tensor.reshape.md

@@ -1,15 +1,21 @@
 # tensor.reshape
 
 ```rust 
-   fn reshape(self: @Tensor<T>, target_shape: Span<usize>) -> Tensor<T>;
+   fn reshape(self: @Tensor<T>, target_shape: Span<i32>, allowzero: bool) -> Tensor<T>;
 ```
 
-Returns a new tensor with the specified target shape and the same data as the input tensor.
+Reshape the input tensor similar to numpy.reshape. First input is the data tensor, second 
+input is a shape tensor which specifies the output shape. It outputs the reshaped tensor. 
+At most one dimension of the new shape can be -1. In this case, the value is inferred from 
+the size of the tensor and the remaining dimensions. A dimension could also be 0, in which case 
+the actual dimension value is unchanged (i.e. taken from the input tensor). If 'allowzero' is set,
+and the new shape includes 0, the dimension will be set explicitly to zero (i.e. not taken from input tensor)
 
 ## Args
 
 * `self`(`@Tensor<T>`) - The input tensor.
-* `target_shape`(Span<usize>) - A span containing the target shape of the tensor.
+* `target_shape`(Span<i32>) - A span containing the target shape of the tensor.
+* `allowzero`(`bool`) - Indicates that if any value in the 'shape' input is set to zero, the zero value is honored, similar to NumPy.
 
 ## Panics
 
@@ -32,7 +38,7 @@ fn reshape_tensor_example() -> Tensor<u32> {
     );
 
     // We can call `reshape` function as follows.
-    return tensor.reshape(target_shape: array![2, 4].span());
+    return tensor.reshape(target_shape: array![2, 4].span(), false);
 }
 >>> [[0,1,2,3], [4,5,6,7]]
 ```

nodegen/file_manager.py

@@ -91,25 +91,36 @@ def base_template(
         This method generates a list of strings that form the template of a Cairo test function,
         including module imports, function definition, and assertions.
         """
-        return [
+        template = [
             *[f"mod input_{i};" for i in range(arg_cnt)],
             *[f"mod output_{i};" for i in range(out_cnt)],
-            *[""],
-            *[""],
+            "",
+            "",
             *[f"use {ref};" for ref in refs],
-            *[""],
-            *["#[test]"],
-            *["#[available_gas(2000000000)]"],
-            *[f"fn test_{name}()" + " {"],
+            "",
+            "#[test]",
+            "#[available_gas(2000000000)]",
+            f"fn test_{name}()" + " {",
             *[f"    let input_{i} = input_{i}::input_{i}();" for i in range(arg_cnt)],
             *[f"    let z_{i} = output_{i}::output_{i}();" for i in range(out_cnt)],
-            *[""],
-            *[f"    let ({', '.join(f'y_{i}' for i in range(out_cnt))}) = {func_sig};"],
-            *[""],
-            *[f"    assert_eq(y_{i}, z_{i});" for i in range(out_cnt)],
-            *["}"],
+            ""
         ]
 
+        # Handling conditional function signature based on the number of outputs
+        if out_cnt > 1:
+            template.append(f"    let ({', '.join(f'y_{i}' for i in range(out_cnt))}) = {func_sig};")
+        else:
+            template.append(f"    let y_0 = {func_sig};")
+
+        # Continue appending to the template
+        template.extend([
+            "",
+            *[f"    assert_eq(y_{i}, z_{i});" for i in range(out_cnt)],
+            "}"
+        ])
+
+        return template
+
     @classmethod
     def sequence_template(cls, name: str, arg_cnt: int, refs: list[str], func_sig: str) -> list[str]:
         """