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

* bug fix

* format

python/mlx/nn/layers/activations.py

@@ -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

python/tests/test_ops.py

@@ -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()