Add utilities to distributions
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
In pytorch, the following is easily possible:
logits = ...
probs = Categorical(logits=logits)
log_prob = probs.log_prob(value)
entropy = probs.entropy()
but when I want to achieve something similar in MLX, I have to manually calculate the log_prob and entropy. Is it possible to add support for these methods as it makes working with distributions in MLX much more convenient (at least for me)
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 reviewing how MLX distributions are currently exposed and compare the PyTorch Categorical example, especially the log_prob and entropy entry points. Define which distributions and inputs the utilities should support, then verify that the resulting API covers the demonstrated workflow and add focused tests for those methods.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100