fix matmul translation

sonos · Sep 11, 2023 · 2ea76c0 · 2ea76c0
1 parent ca45133
commit 2ea76c0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/hir/src/ops/matmul.rs b/hir/src/ops/matmul.rs
@@ -65,7 +65,7 @@ impl Expansion for MatMulInference {
         }
         if implicit_n {
             let b = InOut::In(1);
-            let n_axis = axes.axis((b, axes.rank(b) - 2))?;
+            let n_axis = axes.axis((b, axes.rank(b) - 1))?;
             axes = axes.remove_output_axis(0, n_axis.outputs[0][0])?;
         }
         target.wire_node(