Otherwise I wonder if this really finds too much adoption. You don't want to lockin yourself when you have all those other choices. (Ok, to be fair, as it is discussed here, PyTorch etc might not work optimal yet on Apple Silicon, but I guess this is just a matter of time.)
MLX: An array framework for Apple Silicon
21–30 of 47 posts
Re: MLX: An array framework for Apple Silicon
#22I see that the focus is Apple Silicon, but wouldn't it make sense to also make this available to other hardware (e.g. AMD64 + CUDA) and other platforms (e.g. Linux)? Otherwise I wonder if this really finds too much adoption. You don't want to lockin yourself when you have all those other choices. (Ok, to be fair, as it is discussed here, PyTorch etc might not work optimal yet on Apple Silicon, but I guess this is jus…
I already have this nice, powerful, and expensive hardware sitting on my desk. Why not make the most optimal use of it? I can worry about lock-in once it becomes inadequate for the task.
> PyTorch etc might not work optimal yet on Apple Silicon
So now we can take MLX apart and see how we can use it to improve PyTorch.
Re: MLX: An array framework for Apple Silicon
#23Earlier quoted context omitted.
I think it is a feature from functional programming. In some level, I more agree with jax style.
But you can convert automatically a mutable code into a functional code if that makes things easier. That's what Haskell's `do` notation does, and PyTorch even has `torch.func.functionalize` for that. Immutable should be default, but not compulsory.
arr[1000000] = 1
has to clone the entire array if you want it to be pure, leading to very unpredictable performance. There are also some algorithms that are straight up impossible to (efficiently) implement without mutability. Often, it's exactly those algorithms that are hard to optimize for optimizers like JAX.Specifically in JAX, code that is slow due to copying will often be optimized into mutable code before running for performance reasons. But because JAX still has the gurantees of no mutability, it can do many optimizations such as caching or dead-code elimination.
Re: MLX: An array framework for Apple Silicon
#24I wish there was more information as to how this differs from or improves on Jax. Flax+Jax+OpenXLA seems to finally be building some momentum so when a big player launches yet another competitor, the justification for it would be a good thing to see. What was “not good enough” with Jax? Why did it make sense to put this human time and energy there instead of doubling down on Flax/Jax/OpenXLA? How will this move the n…
Many keep forgeting that CUDA for years is a polyglot platform, C, C++, Fortran, plus anything PTX, some of which also target OpenCL, meaning Haskell, Java, C#, Julia, Futhark, or Python bindings.
Then there are the libraries, and GPGPU graphical debugging tools.
By the way, Modular just announced partnerships with AWS and NVidia for Mojo and related tooling.
Re: MLX: An array framework for Apple Silicon
#25Awni Hannun of Apple AI research team has a deep dive thread on X [0] with some additional info and answers some of the questions [0] https://x.com/awnihannun/status/1732184443451019431?s=46&t=O...
Deep dive? Answering questions? Literally “thumbs up” to “does it use GPU?”? Uh… The links on that thread are the same links from the top of the GitHub repo. I mean, here’s a nitter link https://nitter.net/awnihannun/status/1732184443451019431#m for anyone else who’s interested, but the info on x seems to be a nothing that isn’t already on the GitHub page. The GitHub repo seems to be more useful and informative.
That is how the kidz do deep dives these days.
Re: MLX: An array framework for Apple Silicon
#26I wish there was more information as to how this differs from or improves on Jax. Flax+Jax+OpenXLA seems to finally be building some momentum so when a big player launches yet another competitor, the justification for it would be a good thing to see. What was “not good enough” with Jax? Why did it make sense to put this human time and energy there instead of doubling down on Flax/Jax/OpenXLA? How will this move the n…
Running non-trivial ML workloads on the edge has been on my wishlist for years and it sounds like Apple has just the thing.
Re: MLX: An array framework for Apple Silicon
#27Earlier quoted context omitted.
But you can convert automatically a mutable code into a functional code if that makes things easier. That's what Haskell's `do` notation does, and PyTorch even has `torch.func.functionalize` for that. Immutable should be default, but not compulsory.
Haskell's `do` doesn't allow mutability still, it just allows a syntax that looks a bit more imperative than usual. The problem with all the "convert to pure function" magic is that for example this piece of code arr[1000000] = 1 has to clone the entire array if you want it to be pure, leading to very unpredictable performance. There are also some algorithms that are straight up impossible to (efficiently) implement…
You are correct about in-depth mutations and resulting complications, but that only strengthens my assertion: immutable should be default, but not compulsory (because sometimes you absolutely need them). And mutability doesn't preclude caching or dead-code elimination; you just have to be more careful. Often it's the case that you can convert a mutable code into an immutable form only for the purpose of analysis, which is definitely harder than an immutable code in the first place but not impossible. Scalar compilers have used SSA---an immutable description for mutable programs---for a long time after all.
Re: MLX: An array framework for Apple Silicon
#28I've found the process of porting custom ML models to iOS extremely difficult. AFAIK the only way to leverage Apple Neural Engine (and get the best performance) is to use CoreML. The only documented way to use CoreML is via coremltools, which takes a trace of a PyTorch model and attempts to translate it into a protobuf graph understood by CoreML. This process often fails and requires model changes, or worse "succeeds…
The process becomes particularly frustrating when the model appears to convert successfully, but then fails to produce any output or loses layers entirely. Additionally, the debug information provided by the conversion tool isn't very helpful, adding to the challenge.
As an iOS developer with no prior experience in Python, I found myself in a unique position. I needed to build a custom model for one of my keyboard apps to handle tasks like spellchecking, grammar correction, next-word prediction, and autocompletion. This necessity pushed me to learn Python and PyTorch. After mastering these, I then had to convert my knowledge back to Swift and CoreML.
Ideally, I would have preferred to build my model directly in Swift & CoreML, but the current tools and resources for this approach are limited. This limitation is particularly evident in terms of the ease of use and flexibility that Python and PyTorch offer.
Re: MLX: An array framework for Apple Silicon
#29It seems like it's matching PyTorch's API very closely, which is great. Part of me wishes they took it a step further and just made it completely API-compatible, such that code written for PyTorch could run out-of-the-box with MLX, that would be killer.
Note that there is a Metal backend for PyTorch [0]. Sadly it doesn't work well with codebases that didn't account for it from the start... [0] https://developer.apple.com/metal/pytorch/
There are few limitations left when compared with other backends. Instead of using 'cuda' device, one simply uses 'MPS' as device.
What remains is: the optimizations Pytorch provides (especially compile() with 2.1) focus on cuda and it's historic restrictions that result from CUDA being _not_ unified memory, and lots of energy goes into developing architectural work-arounds in order to limit the copying between graphics HW and CPU memory, resulting in proprietary compilers (like triton) that move parts of the python code into proprietary hardware.
Apple's unified memory would make all of those super complicated architectural workarounds mostly unnecessary (which they demonstrate with their project).
Getting current D/L platforms to support both paradigms (unified/non unified) will be a lot of work. One possible avenue is the MLIR project currently leveraged by Mojo.
Re: MLX: An array framework for Apple Silicon
#30Earlier quoted context omitted.
I'm not sure why there aren't more companies supporting ONNX. It's so nice to use if it's supported by the platform/model.
Not every model can be easily converted to ONNX though, especially with PyTorch.