Awni Hannun
Awni Hannun
I'm not opposed to this. I can mark it as an enhancement. @zcbenz I'm not sure if you are planning to send a PR. Happy to take a look if...
Very impressive perf! Regarding the design, there is a big style difference from other MLX ops which we should change if possible. Basically you do the dispatch at the op-level...
> That makes sense to me, it did feel a little anti-pattern bloating out the graph but the MLX api is just really convenient! Let me give the re-write a...
I might prefer the name `comm` or `comms` over `dist`. I also think `distributed` is fine.. and perhaps better since it is what every other package uses. One can always...
This is very nice and simple! Looks great!
Usage docs coming soon!
- [usage guide](https://ml-explore.github.io/mlx/build/html/usage/distributed.html) - [API spec](https://ml-explore.github.io/mlx/build/html/python/distributed.html)
@kemchenj we have a [custom_vjp transformation](https://github.com/ml-explore/mlx/blob/main/mlx/transforms.h#L181-L190) in C++.
@angeloskath regarding the Python interface. I think in option A is nicer for enabling any transformation to be customized which I think is preferable. I'm not so crazy about the...
This is available in v0.16. See the [documentation for more information](https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.custom_function.html). If you have any feedback or run into any issues please let us know as this is a new...