[Feature] custom extensions

[petertsoi]

Currently, if I want to write a new custom operation and use it in Swift, I would have to maintain my own fork of MLX. It would be great to be able to bring my own CPU/GPU implementations in my own app but still be able to use them like I can in Python.

[awni]

Probably your best bet for this is to use Metal custom kernels which will be added in #137. We also have custom function (which we will plan to add to MLX Swift) to enable custom transformations (like vjp/jvp/vmap). Combining the two you should be able to do everything you can do with a custom extension (with the exception of custom CPU kernels which we don't have yet).