From c4767d110f11f22fca548c5e0801a192cf0e3de6 Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Mon, 27 Oct 2025 11:33:32 -0700 Subject: [PATCH] fix addmm cpu (#2699) --- mlx/backend/cpu/gemms/bnns.cpp | 9 +++++++-- python/tests/test_blas.py | 4 ++-- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/mlx/backend/cpu/gemms/bnns.cpp b/mlx/backend/cpu/gemms/bnns.cpp index 0ba7a9b8..2ec0fd4e 100644 --- a/mlx/backend/cpu/gemms/bnns.cpp +++ b/mlx/backend/cpu/gemms/bnns.cpp @@ -1,5 +1,4 @@ // Copyright © 2023-2024 Apple Inc. - #include #include "mlx/array.h" @@ -49,9 +48,15 @@ void matmul_bnns( size_t K = a_shape[ndim - 1]; BNNSDataType bnns_dtype = to_bnns_dtype(); - #pragma GCC diagnostic push #pragma GCC diagnostic ignored "-Wdeprecated-declarations" + if (beta != 1.0 && beta != 0.0) { + // scale the output + for (auto i = 0; i < batch_size * M * N; ++i) { + out[i] *= beta; + } + beta = 1.0; + } const BNNSLayerParametersBroadcastMatMul gemm_params{ /* float alpha = */ alpha, /* float beta = */ beta, diff --git a/python/tests/test_blas.py b/python/tests/test_blas.py index 82c63e3d..0cc6e621 100644 --- a/python/tests/test_blas.py +++ b/python/tests/test_blas.py @@ -717,8 +717,8 @@ class TestBlas(mlx_tests.MLXTestCase): c = mx.ones((32, 32)).astype(t) a = mx.random.uniform(shape=(32, 32)).astype(t) b = mx.random.uniform(shape=(32, 32)).astype(t) - out = mx.addmm(c, a, b) - expected = a @ b + c + out = mx.addmm(c, a, b, alpha=0.5, beta=2.0) + expected = 0.5 * (a @ b) + 2.0 * c self.assertTrue(mx.allclose(out, expected, rtol=tol, atol=tol)) def test_addmm_grad(self):