Merge pull request #554 from gizatechxyz/develop

Merge Develop into Main

Scarb.toml

@@ -10,8 +10,7 @@ homepage = "https://github.com/gizatechxyz/orion"
 alexandria_merkle_tree = { git = "https://github.com/keep-starknet-strange/alexandria.git", rev = "01a7690" }
 alexandria_data_structures = { git = "https://github.com/keep-starknet-strange/alexandria.git", rev = "01a7690" }
 alexandria_sorting = { git = "https://github.com/keep-starknet-strange/alexandria.git", rev = "01a7690" }
-# TODO: update to https://github.com/influenceth/cubit & change rev
-cubit = { git = "https://github.com/akhercha/cubit.git", rev = "d3869a3" }
+cubit = { git = "https://github.com/influenceth/cubit.git", rev = "6275608" }
 
 [scripts]
 sierra = "cairo-compile . -r"

src/operators/nn/functional/gemm.cairo

@@ -1,3 +1,4 @@
+use alexandria_data_structures::array_ext::SpanTraitExt;
 use core::array::SpanTrait;
 
 use orion::numbers::NumberTrait;
@@ -48,6 +49,14 @@ fn gemm<
 
     match C {
         Option::Some(c) => {
+            let broadcast_c_shape = if c.shape.len() == 1 {
+                array![1].span().concat(c.shape)
+            } else {
+                c.shape
+            };
+
+            let c = Tensor { shape: broadcast_c_shape, data: c.data };
+
             return mul_by_scalar(@A.matmul(@B), alpha) + mul_by_scalar(@c, beta);
         },
         Option::None(_) => { return mul_by_scalar(@A.matmul(@B), alpha); }

src/operators/tensor/helpers.cairo

@@ -496,4 +496,4 @@ impl SpanPartialOrd<T, +Drop<T>, +Copy<T>, +PartialEq<T>, +PartialOrd<T>> of Par
     fn lt(lhs: Span<T>, rhs: Span<T>) -> bool {
         span_cmp(lhs, rhs) < 0
     }
-}
+}