Close a couple edge case bugs: hadamard and addmm on empty inputs (#2177)

* handle hadamard and addmm on empty inputs

* fix
This commit is contained in:
Awni Hannun
2025-05-12 10:48:57 -07:00
committed by GitHub
parent caaa3f1f8c
commit 8f3d208dce
5 changed files with 52 additions and 1 deletions

View File

@@ -589,6 +589,10 @@ class TestBlas(mlx_tests.MLXTestCase):
alpha = 0.5
beta = 2.0
# c must broadcast to the output shape
with self.assertRaises(ValueError):
mx.addmm(mx.zeros((2, 2, 2)), mx.zeros((2, 2)), mx.zeros((2, 2)))
# Regular batched case
a_npy = np.random.normal(0.0, 1.0 / 128, (32, 128, 16)).astype(np.float32)
b_npy = np.random.normal(0.0, 1.0 / 128, (32, 16, 16)).astype(np.float32)
@@ -745,6 +749,19 @@ class TestBlas(mlx_tests.MLXTestCase):
mx.eval(c)
self.assertEqual(c.shape, (0, 0))
c = mx.array(1.0, dtype=mx.float32)
a = mx.array([], dtype=mx.float32)
b = mx.array([], dtype=mx.float32)
out = mx.addmm(c, a, b)
self.assertEqual(out.item(), 1.0)
self.assertEqual(out.shape, ())
a = mx.zeros(shape=(5, 0))
b = mx.zeros(shape=(0, 5))
c = mx.random.uniform(shape=(5, 5))
out = mx.addmm(c, a, b)
self.assertTrue(mx.allclose(out, c))
def test_block_masked_matmul(self):
def ref_block_masked_mm(
a, b, block_size, out_mask=None, lhs_mask=None, rhs_mask=None

View File

@@ -2830,6 +2830,9 @@ class TestOps(mlx_tests.MLXTestCase):
return H
def test_hadamard(self):
with self.assertRaises(ValueError):
mx.hadamard_transform(mx.array([]))
h28_str = """
+------++----++-+--+-+--++--
-+-----+++-----+-+--+-+--++-