Bitwise Inverse (#1862)

* add bitwise inverse

* add vmap + fix nojit

* inverse -> invert

* add to compile + remove unused

python/src/array.cpp

@@ -745,11 +745,10 @@ void init_array(nb::module_& m) {
               throw std::invalid_argument(
                   "Floating point types not allowed with bitwise inversion.");
             }
-            if (a.dtype() != mx::bool_) {
-              throw std::invalid_argument(
-                  "Bitwise inversion not yet supported for integer types.");
+            if (a.dtype() == mx::bool_) {
+              return mx::logical_not(a);
             }
-            return mx::logical_not(a);
+            return mx::bitwise_invert(a);
           })
       .def(
           "__and__",

python/src/ops.cpp

@@ -4833,6 +4833,28 @@ void init_ops(nb::module_& m) {
         Returns:
             array: The bitwise right shift ``a >> b``.
       )pbdoc");
+  m.def(
+      "bitwise_invert",
+      [](const ScalarOrArray& a_, mx::StreamOrDevice s) {
+        auto a = to_array(a_);
+        return mx::bitwise_invert(a, s);
+      },
+      nb::arg(),
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      nb::sig(
+          "def bitwise_invert(a: Union[scalar, array], stream: Union[None, Stream, Device] = None) -> array"),
+      R"pbdoc(
+        Element-wise bitwise inverse.
+
+        Take the bitwise complement of the input.
+
+        Args:
+            a (array): Input array or scalar.
+
+        Returns:
+            array: The bitwise inverse ``~a``.
+      )pbdoc");
   m.def(
       "view",
       [](const ScalarOrArray& a, const mx::Dtype& dtype, mx::StreamOrDevice s) {

python/tests/test_ops.py

@@ -2573,6 +2573,18 @@ class TestOps(mlx_tests.MLXTestCase):
                 out_np = getattr(np, op)(a_np, b_np)
                 self.assertTrue(np.array_equal(np.array(out_mlx), out_np))
 
+        for t in types:
+            a_mlx = a.astype(t)
+            a_np = np.array(a_mlx)
+
+            out_mlx = ~a_mlx
+            out_np = ~a_np
+            self.assertTrue(np.array_equal(np.array(out_mlx), out_np))
+
+            out_mlx = mx.bitwise_invert(a_mlx)
+            out_np = mx.bitwise_invert(a_np)
+            self.assertTrue(np.array_equal(np.array(out_mlx), out_np))
+
         # Check broadcasting
         a = mx.ones((3, 1, 5), dtype=mx.bool_)
         b = mx.zeros((1, 2, 5), dtype=mx.bool_)