Missing something like numpy.bincount

[mateusrodriguesxyz]

While converting a KNN implementation from NumPy to MLX Swift I stumbled upon the usage of a bincount to count the number of occurrences of each value in array. I couldn't find a MLX equivalent so I did the following:

let prediction =  kNearestLabels
        .asArray(Int.self) // MLXArray -> [Int]
        .reduce(into: [:]) { $0[$1, default: 0] += 1 }  // count labels
        .max(by: { $0.1 < $1.1 })!  // max '(label, count)' pair
        .key // label (prediction)

The NumPy code did this in just one line:

np.bincount(k_nearest_labels).argmax()

Is there a more "MLX way" of doing this? If not, maybe MLX should have this bincount method, seems quite handy.

[davidkoski]

Transfer this issue to mlx -- I don't see one in the underlying C++ core. Per the docs

• https://numpy.org/doc/2.1/reference/generated/numpy.bincount.html

it looks like it produces a histogram. I think the problem with the API (as written) is that the output size is variable based on the contents of the input, so it doesn't match the mlx model.

[davidkoski]

Oh, I can't transfer it. Well @awni what do you think?

[awni]

I'll transfer it to MLX. I think we could have a version where the maximum value is provided as an argument.. but without that it would be an op with data-dependent output shape and we don't have support for those in MLX at the moment.

[angeloskath]

Well given a fixed output size N then you can simply write

mx.zeros(N, dtype=mx.int32).at[k_neareset_labels].add(1).argmax()

[mateusrodriguesxyz]

Well given a fixed output size N then you can simply write

mx.zeros(N, dtype=mx.int32).at[k_neareset_labels].add(1).argmax()

Thanks! Can this be done in Swift? I can't find the equivalent to this at and add methods. The best I could do was this:

let a = MLXArray([0, 0, 0])
let idx = MLXArray([0, 1, 0, 1, 1, 2, 1, 2, 2, 2, 2])
for i in idx {
    a[i] = a[i] + 1
}
print(a.argMax())

[davidkoski]

Added an issue to add the missing at() function: ml-explore/mlx-swift#188

[awni]

@angeloskath the one liner you wrote for this is the entire implementation so I’m debating if it’s worth adding this op or closing this. It might be handy to have the op since it’s not entirely obvious.. I’ll mark it as an enhancement for now but I’m also ok to close it.

[angeloskath]

Personally I don't think it is worth an op but I tend to be overly conservative wrt these things ie I don't like ops like relu etc. I think they hide more than they are worth. Otoh jax has it 🤷‍♂️