Bug fix in metal binary kernel dispatch for large arrays (#125)

* bug fix

* format
This commit is contained in:
Awni Hannun
2023-12-10 16:12:31 -08:00
committed by GitHub
parent 0cfbfc9904
commit 71d1fff90a
3 changed files with 10 additions and 4 deletions

View File

@@ -22,7 +22,6 @@ def sigmoid(x):
\text{Sigmoid}(x) = \sigma(x) = \frac{1}{1 + \exp(-x)}
"""
return mx.sigmoid(x)
def relu(x):
@@ -89,10 +88,12 @@ def gelu_fast_approx(x):
"""
return x * mx.sigmoid(1.773 * x)
@_make_activation_module
class Sigmoid(Module):
pass
@_make_activation_module(relu)
class ReLU(Module):
pass

View File

@@ -1305,6 +1305,11 @@ class TestOps(mlx_tests.MLXTestCase):
d_np = np.take(b_mx, np.arange(kth), axis=axis)
self.assertTrue(np.all(d_np <= c_mx))
def test_large_binary(self):
a = mx.ones([1000, 2147484], mx.int8)
b = mx.ones([2147484], mx.int8)
self.assertEqual((a + b)[0, 0].item(), 2)
if __name__ == "__main__":
unittest.main()