Add bincount method to array class

mlx/array.cpp

@@ -1,4 +1,3 @@
-// Copyright © 2023-2024 Apple Inc.
 #include <functional>
 #include <unordered_map>
 
@@ -341,4 +340,22 @@ array::ArrayIterator::reference array::ArrayIterator::operator*() const {
   return reshape(slice(arr, start, end), shape);
 };
 
+array bincount(const array& input, int max_val) {
+  std::unordered_map<int, int> counts;
+  for (int i = 0; i < input.size(); ++i) {
+    int val = input.data<int>()[i];
+    if (val < 0 || val >= max_val) {
+      throw std::out_of_range("Input value out of range");
+    }
+    counts[val]++;
+  }
+
+  std::vector<int> result(max_val, 0);
+  for (const auto& [val, count] : counts) {
+    result[val] = count;
+  }
+
+  return array(result);
+}
+
 } // namespace mlx::core

mlx/array.h

@@ -1,4 +1,3 @@
-// Copyright © 2023 Apple Inc.
 #pragma once
 
 #include <algorithm>
@@ -446,6 +445,9 @@ class array {
 
   ~array();
 
+  // Declare the bincount method in the array class
+  static array bincount(const array& input, int max_val);
+
  private:
   // Initialize the arrays data
   template <typename It>

python/src/array.cpp

@@ -1,4 +1,3 @@
-// Copyright © 2023-2024 Apple Inc.
 #include <cstdint>
 #include <cstring>
 #include <sstream>
@@ -1394,5 +1393,22 @@ void init_array(nb::module_& m) {
           "dtype"_a,
           nb::kw_only(),
           "stream"_a = nb::none(),
-          "See :func:`view`.");
+          "See :func:`view`.")
+      .def_static(
+          "bincount",
+          [](const mx::array& input, int max_val) {
+            return mx::array::bincount(input, max_val);
+          },
+          "input"_a,
+          "max_val"_a,
+          R"pbdoc(
+            Count occurrences of each value in an array.
+
+            Args:
+                input (array): Input array.
+                max_val (int): Maximum value in the input array.
+
+            Returns:
+                array: Array with the counts of each value in the input array.
+          )pbdoc");
 }

python/tests/test_array.py

@@ -1994,6 +1994,39 @@ def t():
         used = get_mem()
         self.assertEqual(expected, used)
 
+    def test_bincount(self):
+        # Test case 1: Basic functionality
+        input_array = mx.array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3])
+        max_val = 4
+        expected_output = mx.array([1, 2, 3, 4])
+        self.assertTrue(mx.array_equal(mx.array.bincount(input_array, max_val), expected_output))
+
+        # Test case 2: Input with zeros
+        input_array = mx.array([0, 0, 0, 1, 1, 2])
+        max_val = 3
+        expected_output = mx.array([3, 2, 1])
+        self.assertTrue(mx.array_equal(mx.array.bincount(input_array, max_val), expected_output))
+
+        # Test case 3: Input with negative values (should raise an error)
+        input_array = mx.array([0, -1, 1, 2])
+        max_val = 3
+        with self.assertRaises(ValueError):
+            mx.array.bincount(input_array, max_val)
 
-if __name__ == "__main__":
-    unittest.main()
+        # Test case 4: Input with values out of range (should raise an error)
+        input_array = mx.array([0, 1, 2, 3, 4])
+        max_val = 3
+        with self.assertRaises(ValueError):
+            mx.array.bincount(input_array, max_val)
+
+        # Test case 5: Empty input array
+        input_array = mx.array([])
+        max_val = 3
+        expected_output = mx.array([0, 0, 0])
+        self.assertTrue(mx.array_equal(mx.array.bincount(input_array, max_val), expected_output))
+
+        # Test case 6: Input with all values the same
+        input_array = mx.array([2, 2, 2, 2])
+        max_val = 3
+        expected_output = mx.array([0, 0, 4])
+        self.assertTrue(mx.array_equal(mx.array.bincount(input_array, max_val), expected_output))