mirror of
https://github.com/ml-explore/mlx.git
synced 2025-10-18 15:28:16 +08:00
Bug fix in metal binary kernel dispatch for large arrays (#125)
* bug fix * format
This commit is contained in:
@@ -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
|
||||
|
@@ -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()
|
||||
|
Reference in New Issue
Block a user