[Feature] More accurate reductions for low precision types
Open
Nobody has claimed this yet.
enhancement
low priority
- Dominant language
- C++
- Stars
- 28.5k
- Forks
- 2.3k
- Avg merge
- 3d 8h
- Merged PRs (30d)
- 62
Description
Our reductions are quite naive and can be less accurate particularly in lower precision (mx.float16).
NumPy and PyTorch (MPS) seem to use more sophisticated reductions. E.g. for NumPy:
For floating point numbers the numerical precision of sum (and np.add.reduce) is in general limited by directly adding each number individually to the result causing rounding errors in every step. However, often numpy will use a numerically better approach (partial pairwise summation) leading to improved precision in many use-cases. This improved precision is always provided when no axis is given.
See #483 for a little more discussion.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reading the discussion in #483 and locating MLX's reduction implementations for low-precision types. Compare their behavior with the NumPy and PyTorch MPS approaches described here. The work is done when reductions such as sum provide measurably improved accuracy for types including mx.float16, with appropriate validation against the intended behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100