From 293114a31416797d19877ab7db211b7d4e7176fb Mon Sep 17 00:00:00 2001
From: Mathieu Poumeyrol <kali@zoy.org>
Date: Tue, 3 Sep 2024 10:10:35 +0200
Subject: [PATCH] unchecked axis access

---
 core/src/ops/matmul/optimized.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/core/src/ops/matmul/optimized.rs b/core/src/ops/matmul/optimized.rs
index 9f398c8076..312019cf29 100644
--- a/core/src/ops/matmul/optimized.rs
+++ b/core/src/ops/matmul/optimized.rs
@@ -291,8 +291,8 @@ impl EvalOp for OptMatMul {
                 let uops: Vec<FusedSpec> =
                     self.micro_ops.iter().map(|o| o.resolve_trivial(&inputs, &mut c)).collect();
                 self.mmm.run_with_scratch_space(
-                    self.m.to_i64()? as usize,
-                    self.n.to_i64()? as usize,
+                    *c_shape.get_unchecked(self.c_m_axis),
+                    *c_shape.get_unchecked(self.c_n_axis),
                     scratch.as_mut(),
                     &uops,
                 )?;