Masked gather

[CC-Yeh]

It would be great to support boolean-mask indexing, e.g.:

out = values[mask]

to gather all elements where mask == True.

Right now, this requires manual index construction and workarounds (e.g., prefix + scatter_add), which is both verbose and slower than a dedicated masked gather. A built-in op for boolean mask selection would simplify code and improve performance for common use cases where we extract values based on a mask.

[zcbenz]

Is this the same with https://github.com/ml-explore/mlx/issues/246?

[CC-Yeh]

Yes, I think masked_gather would be a nice way to support #246.