Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add bincount method to array class #1874

Closed
wants to merge 1 commit into from

Conversation

anupamme
Copy link

@anupamme anupamme commented Feb 17, 2025

Fixes #1786

Add bincount method to MLX to count occurrences of each value in an array.

  • mlx/array.cpp: Implement bincount method using std::unordered_map to store counts and return an array with the counts of each value in the input array.
  • mlx/array.h: Declare bincount method in the array class.
  • python/src/array.cpp: Add bincount method to the Python API for the array class and implement it using the bincount method from mlx/array.cpp.
  • python/tests/test_array.py: Add test cases for the bincount method in the array class to verify its correctness with various input arrays.

For more details, open the Copilot Workspace session.

Fixes ml-explore#1786

Add `bincount` method to MLX to count occurrences of each value in an array.

* **mlx/array.cpp**: Implement `bincount` method using `std::unordered_map` to store counts and return an array with the counts of each value in the input array.
* **mlx/array.h**: Declare `bincount` method in the `array` class.
* **python/src/array.cpp**: Add `bincount` method to the Python API for the `array` class and implement it using the `bincount` method from `mlx/array.cpp`.
* **python/tests/test_array.py**: Add test cases for the `bincount` method in the `array` class to verify its correctness with various input arrays.

---

For more details, open the [Copilot Workspace session](https://copilot-workspace.githubnext.com/ml-explore/mlx/issues/1786?shareId=XXXX-XXXX-XXXX-XXXX).
@angeloskath
Copy link
Member

Hi, the linked issue has a much more concise and correct implementation as follows

def bincount(input, max_val: int):
    return mx.zeros(max_val, dtype=mx.int32).at[input].add(1)

this also runs in the GPU and doesn't require evaluating input before computing the bin count.

@anupamme anupamme deleted the add-bincount branch February 17, 2025 09:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Missing something like numpy.bincount
2 participants