Merge branch 'develop' into pr/591

.github/workflows/test.yaml

@@ -9,5 +9,5 @@ jobs:
       - uses: actions/checkout@v3
       - uses: software-mansion/setup-scarb@v1
         with:
-          scarb-version: "2.5.3"
+          scarb-version: "2.6.4"
       - run: scarb test --workspace && scarb fmt --workspace

.tool-versions

@@ -1 +1 @@
-scarb 2.5.3
+scarb 2.6.4

Scarb.toml

@@ -1,6 +1,6 @@
 [package]
 name = "orion"
-version = "0.2.4"
+version = "0.2.5"
 cairo-version = "2.5.3"
 edition = "2023_10"
 description = "ONNX Runtime in Cairo for verifiable ML inference using STARK"

docgen/src/main.rs

@@ -59,6 +59,14 @@ fn main() {
     doc_trait(trait_path, doc_path, label);
     doc_functions(trait_path, doc_path, trait_name, label);
 
+    // TREE ENSEMBLE DOC
+    let trait_path = "src/operators/ml/tree_ensemble/tree_ensemble.cairo";
+    let doc_path = "docs/framework/operators/machine-learning/tree-ensemble";
+    let label = "tree_ensemble";
+    let trait_name: &str = "TreeEnsembleTrait";
+    doc_trait(trait_path, doc_path, label);
+    doc_functions(trait_path, doc_path, trait_name, label);
+
     // LINEAR REGRESSOR DOC
     let trait_path = "src/operators/ml/linear/linear_regressor.cairo";
     let doc_path = "docs/framework/operators/machine-learning/linear-regressor";

docs/SUMMARY.md

@@ -179,6 +179,7 @@
     * [nn.conv](framework/operators/neural-network/nn.conv.md)
     * [nn.depth_to_space](framework/operators/neural-network/nn.depth_to_space.md)
     * [nn.space_to_depth](framework/operators/neural-network/nn.space_to_depth.md)
+    * [nn.max\_pool](framework/operators/neural-network/nn.max\_pool.md)
   * [Machine Learning](framework/operators/machine-learning/README.md)
     * [Tree Ensemble Classifier](framework/operators/machine-learning/tree-ensemble-classifier/README.md)
       * [tree\_ensemble\_classifier.predict](framework/operators/machine-learning/tree-ensemble-classifier/tree\_ensemble\_classifier.predict.md)

docs/framework/compatibility.md

@@ -47,6 +47,7 @@ You can see below the list of current supported ONNX Operators:
 |            [Col2im](operators/neural-network/nn.col2im\_sigmoid.md)           | :white\_check\_mark: |
 |       [ConvTranspose](operators/neural-network/nn.conv\_transpose_.md)        | :white\_check\_mark: |
 |                  [Conv](operators/neural-network/nn.conv.md)                  | :white\_check\_mark: |
+|              [MaxPool](operators/neural-network/nn.max\_pool.md)              | :white\_check\_mark: |
 |                    [Sinh](operators/tensor/tensor.sinh.md)                    | :white\_check\_mark: |
 |                   [Asinh](operators/tensor/tensor.asinh.md)                   | :white\_check\_mark: |
 |                   [Atanh](operators/tensor/tensor.atanh.md)                   | :white\_check\_mark: |

docs/framework/operators/machine-learning/tree-ensemble/README.md

@@ -0,0 +1,22 @@
+# Tree Ensemble 
+
+`TreeEnsembleTrait` provides a trait definition for tree ensemble problem.
+
+```rust
+use orion::operators::ml::TreeEnsembleTrait;
+```
+
+### Data types
+
+Orion supports currently only fixed point data types for `TreeEnsembleTrait`.
+
+| Data type            | dtype                                                         |
+| -------------------- | ------------------------------------------------------------- |
+| Fixed point (signed) | `TreeEnsembleTrait<FP8x23 \| FP16x16 \| FP64x64 \| FP32x32>` |
+
+
+***
+
+| function | description |
+| --- | --- |
+| [`tree_ensemble.predict`](tree_ensemble.predict.md) | Returns the regressed values for each input in a batch. |

docs/framework/operators/machine-learning/tree-ensemble/tree_ensemble.predict.md

@@ -0,0 +1,139 @@
+# TreeEnsemble::predict
+
+```rust 
+   fn predict(X: @Tensor<T>,
+                  nodes_splits: Tensor<T>,
+                  nodes_featureids: Span<usize>,
+                  nodes_modes: Span<MODE>,
+                  nodes_truenodeids: Span<usize>,
+                  nodes_falsenodeids: Span<usize>,
+                  nodes_trueleafs: Span<usize>,
+                  nodes_falseleafs: Span<usize>,
+                  leaf_targetids: Span<usize>,
+                  leaf_weights: Tensor<T>,
+                  tree_roots: Span<usize>,
+                  post_transform: POST_TRANSFORM,
+                  aggregate_function: AGGREGATE_FUNCTION,
+                  nodes_hitrates: Option<Tensor<T>>,
+                  nodes_missing_value_tracks_true: Option<Span<usize>>,
+                  membership_values: Option<Tensor<T>>,
+                  n_targets: usize
+              ) -> MutMatrix::<T>;
+```
+
+Tree Ensemble operator. Returns the regressed values for each input in a batch. Inputs have dimensions [N, F] where N is the input batch size and F is the number of input features. Outputs have dimensions [N, num_targets] where N is the batch size and num_targets is the number of targets, which is a configurable attribute.
+
+## Args
+
+* `X`:  Input 2D tensor.
+* `nodes_splits`: Thresholds to do the splitting on for each node with mode that is not 'BRANCH_MEMBER'.
+* `nodes_featureids`: Feature id for each node.
+* `nodes_modes`: The comparison operation performed by the node. This is encoded as an enumeration of 'NODE_MODE::LEQ', 'NODE_MODE::LT', 'NODE_MODE::GTE', 'NODE_MODE::GT', 'NODE_MODE::EQ', 'NODE_MODE::NEQ', and 'NODE_MODE::MEMBER'
+* `nodes_truenodeids`: If `nodes_trueleafs` is 0 (false) at an entry, this represents the position of the true branch node. 
+* `nodes_falsenodeids`: If `nodes_falseleafs` is 0 (false) at an entry, this represents the position of the false branch node.
+* `nodes_trueleafs`: 1 if true branch is leaf for each node and 0 an interior node.
+* `nodes_falseleafs`: 1 if true branch is leaf for each node and 0 an interior node.
+* `leaf_targetids`: The index of the target that this leaf contributes to (this must be in range `[0, n_targets)`).
+* `leaf_weights`: The weight for each leaf.
+* `tree_roots`: Index into `nodes_*` for the root of each tree. The tree structure is derived from the branching of each node.
+* `post_transform`: Indicates the transform to apply to the score.One of 'POST_TRANSFORM::NONE', 'POST_TRANSFORM::SOFTMAX', 'POST_TRANSFORM::LOGISTIC', 'POST_TRANSFORM::SOFTMAX_ZERO'  or 'POST_TRANSFORM::PROBIT' ,
+* `aggregate_function`: Defines how to aggregate leaf values within a target. One of 'AGGREGATE_FUNCTION::AVERAGE', 'AGGREGATE_FUNCTION::SUM', 'AGGREGATE_FUNCTION::MIN', 'AGGREGATE_FUNCTION::MAX` defaults to 'AGGREGATE_FUNCTION::SUM' 
+* `nodes_hitrates`: Popularity of each node, used for performance and may be omitted.
+* `nodes_missing_value_tracks_true`: For each node, define whether to follow the true branch (if attribute value is 1) or false branch (if attribute value is 0) in the presence of a NaN input feature. This attribute may be left undefined and the default value is false (0) for all nodes.
+* `membership_values`: Members to test membership of for each set membership node. List all of the members to test again in the order that the 'BRANCH_MEMBER' mode appears in `node_modes`, delimited by `NaN`s. Will have the same number of sets of values as nodes with mode 'BRANCH_MEMBER'. This may be omitted if the node doesn't contain any 'BRANCH_MEMBER' nodes.
+* `n_targets`: The total number of targets.
+
+
+## Returns
+
+* Output of shape [Batch Size, Number of targets]
+
+## Type Constraints
+
+`TreeEnsembleClassifier` and `X` must be fixed points
+
+## Examples
+
+```rust
+use orion::numbers::FP16x16;
+use orion::operators::tensor::{Tensor, TensorTrait, FP16x16Tensor, U32Tensor};
+use orion::operators::ml::{TreeEnsembleTrait,POST_TRANSFORM, AGGREGATE_FUNCTION, NODE_MODE};
+use orion::operators::matrix::{MutMatrix, MutMatrixImpl};
+use orion::numbers::NumberTrait;
+
+fn example_tree_ensemble_one_tree() ->  MutMatrix::<FP16x16> {
+    let mut shape = ArrayTrait::<usize>::new();
+    shape.append(3);
+    shape.append(2);
+
+    let mut data = ArrayTrait::new();
+    data.append(FP16x16 { mag: 78643, sign: false });
+    data.append(FP16x16 { mag: 222822, sign: false });
+    data.append(FP16x16 { mag: 7864, sign: true });
+    data.append(FP16x16 { mag: 108789, sign: false });
+    data.append(FP16x16 { mag: 271319, sign: false });
+    data.append(FP16x16 { mag: 115998, sign: false });
+    let mut X = TensorTrait::new(shape.span(), data.span());
+
+    let mut shape = ArrayTrait::<usize>::new();
+    shape.append(4);
+
+    let mut data = ArrayTrait::new();
+    data.append(FP16x16 { mag: 342753, sign: false });
+    data.append(FP16x16 { mag: 794296, sign: false });
+    data.append(FP16x16 { mag: 801505, sign: true });
+    data.append(FP16x16 { mag: 472514, sign: false });
+    let leaf_weights = TensorTrait::new(shape.span(), data.span());
+
+    let mut shape = ArrayTrait::<usize>::new();
+    shape.append(3);
+
+    let mut data = ArrayTrait::new();
+    data.append(FP16x16 { mag: 205783, sign: false });
+    data.append(FP16x16 { mag: 78643, sign: false });
+    data.append(FP16x16 { mag: 275251, sign: false });
+    let nodes_splits = TensorTrait::new(shape.span(), data.span());
+
+    let membership_values = Option::None;
+
+    let n_targets = 2;
+    let aggregate_function = AGGREGATE_FUNCTION::SUM;
+    let nodes_missing_value_tracks_true = Option::None;
+    let nodes_hitrates = Option::None;
+    let post_transform = POST_TRANSFORM::NONE;
+
+    let tree_roots: Span<usize> = array![0].span();
+    let nodes_modes: Span<MODE> = array![MODE::LEQ, MODE::LEQ, MODE::LEQ].span();
+
+    let nodes_featureids: Span<usize> = array![0, 0, 0].span();
+    let nodes_truenodeids: Span<usize> = array![1, 0, 1].span();
+    let nodes_trueleafs: Span<usize> = array![0, 1, 1].span();
+    let nodes_falsenodeids: Span<usize> = array![2, 2, 3].span();
+    let nodes_falseleafs: Span<usize> = array![0, 1, 1].span();
+    let leaf_targetids: Span<usize> = array![0, 1, 0, 1].span();
+
+    return TreeEnsembleTrait::predict(
+        @X,
+        nodes_splits,
+        nodes_featureids,
+        nodes_modes,
+        nodes_truenodeids,
+        nodes_falsenodeids,
+        nodes_trueleafs,
+        nodes_falseleafs,
+        leaf_targetids,
+        leaf_weights,
+        tree_roots,
+        post_transform,
+        aggregate_function,
+        nodes_hitrates,
+        nodes_missing_value_tracks_true,
+        membership_values,
+        n_targets
+    );
+}
+
+>>> [[ 5.23  0.  ]
+     [ 5.23  0.  ]
+     [ 0.   12.12]]    
+```

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

@@ -0,0 +1,107 @@
+
+# NNTrait::max_pool
+
+```rust
+    fn max_pool(
+    X: @Tensor<T>,
+    auto_pad: Option<AUTO_PAD>,
+    ceil_mode: Option<usize>,
+    dilations: Option<Span<usize>>,
+    kernel_shape: Span<usize>,
+    pads: Option<Span<usize>>,
+    storage_order: Option<usize>,
+    strides: Option<Span<usize>>,
+    output_len: usize,
+) -> (Tensor<T>, Option<Tensor<usize>>);
+```
+
+MaxPool consumes an input tensor X and applies max pooling across the tensor according to kernel sizes, stride sizes, and pad lengths. max pooling consisting of computing the max on all values of a subset of the input tensor according to the kernel size and downsampling the data into the output tensor Y for further processing. The output spatial shape is calculated differently depending on whether explicit padding is used, where pads is employed, or auto padding is used, where auto_pad is utilized.
+
+## Args
+
+* `X`(`@Tensor<T>`) - Input data tensor from the previous operator; dimensions for image case are (N x C x H x W), where N is the batch size, C is the number of channels, and H and W are the height and the width of the data. For non image case, the dimensions are in the form of (N x C x D1 x D2 ... Dn), where N is the batch size. 
+* `auto_pad`(`Option<AUTO_PAD>`) - Default is NOTSET, auto_pad must be either NOTSET, SAME_UPPER, SAME_LOWER or VALID. NOTSET means explicit padding is used. SAME_UPPER or SAME_LOWER mean pad the input so that `output_shape[i] = ceil(input_shape[i] / strides[i])` for each axis `i`.
+* `ceil_mode`(`Option<usize>`) - Default is 1, Whether to use ceil or floor (default) to compute the output shape.
+* `dilations`(`Option<Span<usize>>`) - Dilation value along each spatial axis of the filter. If not present, the dilation defaults to 1 along each spatial axis.
+* `kernel_shape`(`Span<usize>`) - The size of the kernel along each axis.
+* `pads`(`Option<Span<usize>>`) - Padding for the beginning and ending along each spatial axis, it can take any value greater than or equal to 0. The value represent the number of pixels added to the beginning and end part of the corresponding axis. `pads` format should be as follow [x1_begin, x2_begin...x1_end, x2_end,...], where xi_begin the number of pixels added at the beginning of axis `i` and xi_end, the number of pixels added at the end of axis `i`. This attribute cannot be used simultaneously with auto_pad attribute. If not present, the padding defaults to 0 along start and end of each spatial axis.
+* `storage_order`(`Option<usize>`) - Default is 0, The storage order of the tensor. 0 is row major, and 1 is column major. 
+* `strides`(`Option<Span<usize>>`) - Stride along each spatial axis. If not present, the stride defaults to 1 along each spatial axis.
+* `output_len`(`Option<usize>`) - Default is 1, If set to 2, return the indices tensor.
+
+## Returns
+
+A `Tensor<T>` that contains the result of the max pool.
+A `Option<Tensor<usize>>` with the indices tensor from max pooling across the input tensor. The dimensions of indices are the same as output tensor. 
+## Examples
+
+```rust
+use orion::operators::nn::NNTrait;
+use orion::numbers::FixedTrait;
+use orion::operators::nn::FP16x16NN;
+use orion::numbers::FP16x16;
+use orion::operators::tensor::{Tensor, TensorTrait, FP16x16Tensor};
+
+
+fn example_max_pool() -> (Tensor<FP16x16>, Option<Tensor<usize>>) {
+    let mut shape = ArrayTrait::<usize>::new();
+   shape.append(1);
+   shape.append(1);
+   shape.append(5);
+   shape.append(5);
+   let mut data = ArrayTrait::new();
+   data.append(FP16x16 { mag: 65536, sign: false });
+   data.append(FP16x16 { mag: 131072, sign: false });
+   data.append(FP16x16 { mag: 196608, sign: false });
+   data.append(FP16x16 { mag: 262144, sign: false });
+   data.append(FP16x16 { mag: 327680, sign: false });
+   data.append(FP16x16 { mag: 393216, sign: false });
+   data.append(FP16x16 { mag: 458752, sign: false });
+   data.append(FP16x16 { mag: 524288, sign: false });
+   data.append(FP16x16 { mag: 589824, sign: false });
+   data.append(FP16x16 { mag: 655360, sign: false });
+   data.append(FP16x16 { mag: 720896, sign: false });
+   data.append(FP16x16 { mag: 786432, sign: false });
+   data.append(FP16x16 { mag: 851968, sign: false });
+   data.append(FP16x16 { mag: 917504, sign: false });
+   data.append(FP16x16 { mag: 983040, sign: false });
+   data.append(FP16x16 { mag: 1048576, sign: false });
+   data.append(FP16x16 { mag: 1114112, sign: false });
+   data.append(FP16x16 { mag: 1179648, sign: false });
+   data.append(FP16x16 { mag: 1245184, sign: false });
+   data.append(FP16x16 { mag: 1310720, sign: false });
+   data.append(FP16x16 { mag: 1376256, sign: false });
+   data.append(FP16x16 { mag: 1441792, sign: false });
+   data.append(FP16x16 { mag: 1507328, sign: false });
+   data.append(FP16x16 { mag: 1572864, sign: false });
+   data.append(FP16x16 { mag: 1638400, sign: false });
+   let mut X = TensorTrait::new(shape.span(), data.span());
+   return NNTrait::max_pool(
+       @X,
+       Option::None,
+       Option::None,
+       Option::None,
+       array![5, 5, 5].span(),
+       Option::Some(array![2, 2, 2, 2].span()),
+       Option::None,
+       Option::None,
+       1
+   );
+
+}
+
+>>> ([
+           [
+               [
+                   [13, 14, 15, 15, 15],
+                   [18, 19, 20, 20, 20],
+                   [23, 24, 25, 25, 25],
+                   [23, 24, 25, 25, 25],
+                   [23, 24, 25, 25, 25],
+               ]
+           ]
+       ], 
+       Option::None)
+
+
+````

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.