Merge branch 'develop' into pr/595

gizatechxyz · Apr 22, 2024 · c67cb4a · c67cb4a
2 parents 6ab6d8d + 2a2c8e5
commit c67cb4a
Show file tree

Hide file tree

Showing 619 changed files with 21,429 additions and 11,807 deletions.
diff --git a/.DS_Store b/.DS_Store
diff --git a/.github/workflows/test.yaml b/.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
diff --git a/.tool-versions b/.tool-versions
@@ -1 +1 @@
-scarb 2.5.3
+scarb 2.6.4
diff --git a/Scarb.toml b/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"

diff --git a/docgen/src/main.rs b/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";

diff --git a/docs/SUMMARY.md b/docs/SUMMARY.md
@@ -113,6 +113,7 @@
     * [tensor.qlinear\_matmul](framework/operators/tensor/tensor.qlinear\_matmul.md)
     * [tensor.qlinear\_concat](framework/operators/tensor/tensor.qlinear\_concat.md)
     * [tensor.qlinear\_leakyrelu](framework/operators/tensor/tensor.qlinear\_leakyrelu.md)
+    * [tensor.qlinear\_conv](framework/operators/tensor/tensor.qlinear\_conv.md)
     * [tensor.nonzero](framework/operators/tensor/tensor.nonzero.md)
     * [tensor.squeeze](framework/operators/tensor/tensor.squeeze.md)
     * [tensor.unsqueeze](framework/operators/tensor/tensor.unsqueeze.md)
@@ -174,11 +175,13 @@
     * [nn.gemm](framework/operators/neural-network/nn.gemm.md)
     * [nn.grid\_sample](framework/operators/neural-network/nn.grid\_sample.md)
     * [nn.col2im](framework/operators/neural-network/nn.col2im.md)
-    * [nn.conv_transpose](framework/operators/neural-network/nn.conv\_transpose.md)
+    * [nn.conv\_transpose](framework/operators/neural-network/nn.conv\_transpose.md)
     * [nn.conv](framework/operators/neural-network/nn.conv.md)
     * [nn.conv_integer](framework/operators/neural-network/nn.conv\_integer.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)
+    * [nn.deform\_conv](framework/operators/neural-network/nn.deform\_conv_.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)

diff --git a/docs/framework/compatibility.md b/docs/framework/compatibility.md
@@ -48,6 +48,8 @@ You can see below the list of current supported ONNX Operators:
 |       [ConvTranspose](operators/neural-network/nn.conv\_transpose_.md)        | :white\_check\_mark: |
 |                  [Conv](operators/neural-network/nn.conv.md)                  | :white\_check\_mark: |
 |         [ConvInteger](operators/neural-network/nn.conv\_integer_.md)          | :white\_check\_mark: |
+|              [MaxPool](operators/neural-network/nn.max\_pool.md)              | :white\_check\_mark: |
+|        [DeformConv](operators/neural-network/nn.deform\_conv_.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: |
@@ -67,6 +69,7 @@ You can see below the list of current supported ONNX Operators:
 |             [QlinearAdd](operators/tensor/tensor.qlinear\_add.md)             | :white\_check\_mark: |
 |             [QlinearMul](operators/tensor/tensor.qlinear\_mul.md)             | :white\_check\_mark: |
 |       [QLinearLeakyRelu](operators/tensor/tensor.qlinear\_leakyrelu.md)       | :white\_check\_mark: |
+|            [QLinearConv](operators/tensor/tensor.qlinear\_conv_.md)           | :white\_check\_mark: |
 |                 [Nonzero](operators/tensor/tensor.nonzero.md)                 | :white\_check\_mark: |
 |                 [Squeeze](operators/tensor/tensor.squeeze.md)                 | :white\_check\_mark: |
 |               [Unsqueeze](operators/tensor/tensor.unsqueeze.md)               | :white\_check\_mark: |

diff --git a/docs/framework/operators/machine-learning/tree-ensemble/README.md b/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. |
diff --git a/docs/framework/operators/machine-learning/tree-ensemble/tree_ensemble.predict.md b/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]]    
+```
diff --git a/docs/framework/operators/neural-network/nn.deform_conv.md b/docs/framework/operators/neural-network/nn.deform_conv.md
@@ -0,0 +1,152 @@
+# NNTrait::deform_conv
+
+```rust
+   fn deform_conv(
+       X: @Tensor<T>,
+       W: @Tensor<T>,
+       offset: @Tensor<T>,
+       B: Option<Span<T>>,
+       mask: Option<Tensor<T>>,
+       dilations: Option<Span<usize>>,
+       group: Option<usize>,
+       kernel_shape: Option<Span<usize>>,
+       offset_group: Option<usize>,
+       pads: Option<Span<usize>>,
+       strides: Option<Span<usize>>,
+   ) -> Tensor<T>
+```
+
+Performs deformable convolution as described in https://arxiv.org/abs/1703.06211 and https://arxiv.org/abs/1811.11168. This operator specification supports the 2-D case.
+
+## Args
+
+   X: @Tensor<T>,
+   W: @Tensor<T>,
+   offset: @Tensor<T>,
+   B: Option<Span<T>>,
+   mask: Option<Tensor<T>>,
+   dilations: Option<Span<usize>>,
+   group: Option<usize>,
+   kernel_shape: Option<Span<usize>>,
+   offset_group: Option<usize>,
+   pads: Option<Span<usize>>,
+   strides: Option<Span<usize>>,
+
+* `X`(`@Tensor<T>`) - Input data tensor. For 2D image data, it has shape (N, C, H, W) where N is the batch size, C is the number of input channels, and H and W are the height and width. 
+* `W`(`@Tensor<T>`) - Weight tensor that will be used in the convolutions. It has shape (oC, C/group, kH, kW), where oC is the number of output channels and kH and kW are the kernel height and width.
+* `offset`(`@Tensor<T>`) - Offset tensor denoting the offset for the sampling locations in the convolution kernel. It has shape (N, offset_group * kH * kW * 2, oH, oW) for 2D data
+* `B`(`Option<Span<T>>`) - Default is a tensor of zeros, optional 1D bias of length oC to be added to the convolution.
+* `mask`(`Option<Tensor<T>>`) -  Default is a tensor of ones, the mask tensor to be applied to each position in the convolution kernel. It has shape (N, offset_group * kH * kW, oH, oW) for 2D data.
+* `dilations`(`Option<Span<usize>>`) - Default is 1 along each axis, dilation value along each spatial axis of the kernel.
+* `group`(`usize`) - Default is 1, number of groups the input and output channels, C and oC, are divided into.
+* `kernel_shape`(`Option<Span<usize>>`) - Shape of the convolution kernel. If not present, it is inferred from the shape of input W.
+* `offset_group`(`Option<usize>`) - Default is 1, number of groups of offset. C must be divisible by offset_group.
+* `pads`(`Option<Span<usize>>`) - Default is 0 along each axis, padding for the beginning and end along each spatial axis. The values represent the number of pixels added to the beginning and end of the corresponding axis and can take any nonnegative value.
+* `strides`(`Option<Span<usize>>`) - Default is 1 along each axis, stride along each spatial axis.
+
+## Returns
+
+A `Tensor<T>` output tensor that contains the result of convolution.
+
+## Examples
+
+```rust
+fn example_deform_conv() -> Tensor<FP16x16> {
+    let mut shape = ArrayTrait::<usize>::new();
+    shape.append(1);
+    shape.append(1);
+    shape.append(3);
+    shape.append(3);
+
+    let mut data = ArrayTrait::new();
+    data.append(FP16x16 { mag: 0, sign: false });
+    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 });
+    let mut X = TensorTrait::new(shape.span(), data.span());
+
+    let mut shape = ArrayTrait::<usize>::new();
+    shape.append(1);
+    shape.append(1);
+    shape.append(2);
+    shape.append(2);
+
+    let mut data = ArrayTrait::new();
+    data.append(FP16x16 { mag: 65536, sign: false });
+    data.append(FP16x16 { mag: 65536, sign: false });
+    data.append(FP16x16 { mag: 65536, sign: false });
+    data.append(FP16x16 { mag: 65536, sign: false });
+    let mut W = TensorTrait::new(shape.span(), data.span());
+
+    let mut shape = ArrayTrait::<usize>::new();
+    shape.append(1);
+    shape.append(8);
+    shape.append(2);
+    shape.append(2);
+
+    let mut data = ArrayTrait::new();
+    data.append(FP16x16 { mag: 32768, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 6553, sign: true });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    data.append(FP16x16 { mag: 0, sign: false });
+    let mut offset = TensorTrait::new(shape.span(), data.span());
+
+
+    return NNTrait::deform_conv(
+        @X,
+        @W,
+        @offset,
+        Option::None,
+        Option::None,
+        Option::None,
+        Option::None,
+        Option::Some(array![2, 2].span()),
+        Option::None,
+        Option::Some(array![0, 0, 0, 0].span()),
+        Option::None,
+    );
+}
+
+>>> [
+        [
+            [
+                [9.5, 11.9],  
+                [20.0, 24.0],
+            ]
+        ]
+    ]
+
+````