Fixed computed output shape of matmul op.

webmachinelearning · Nov 23, 2022 · fe6aff2 · fe6aff2
1 parent 7f42089
commit fe6aff2
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/src/nn/ops/binary.ts b/src/nn/ops/binary.ts
@@ -31,11 +31,11 @@ export abstract class Binary extends SingleOutputOperation {
         const rankA = a.rank;
         const rankB = b.rank;
         if (rankA === 1 && rankB === 1) {
-          outputShape = [];
-        } else if (rankA === 2 && rankB === 1) {
-          outputShape = [a.shape[0], 1];
-        } else if (rankA === 1 && rankB === 2) {
-          outputShape = [1, b.shape[1]];
+          outputShape = []; // scalar
+        } else if (rankA >= 2 && rankB === 1) {
+          outputShape = [...a.shape.slice(0, -1), 1];
+        } else if (rankA === 1 && rankB >= 2) {
+          outputShape = [...b.shape.slice(0, -2), 1, b.shape[1]];
         } else if (rankA >= 2 && rankB >= 2) {
           outputShape = utils.getBroadcastShape(a.shape.slice(0, -2),
               b.shape.slice(0, -2));