Hacker News Reader

Comment by kiratp | Hacker News Reader

I 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 needle at all against Nvidia?

I guess the fact that Google is pushing OpenXLA is making the other giants not want to truly lean in?

I don’t want 10 competing “choices”. I want one clear, open, competitor to Cuda that works on all the competing hardware.

View on HN

The only way to beat CUDA like in many other API cases, is by middleware.

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.

The fact that I can run this on mobile (iOS) devices using the C++ interface makes all the difference for me. I find that extremely refreshing among all those other Python/Server/PC -only frameworks.

Running non-trivial ML workloads on the edge has been on my wishlist for years and it sounds like Apple has just the thing.

I’m somewhat concerned that it doesn’t support ANE from the get-go (they list CPU and GPU, for now). Not even for inference only.

Other frameworks have (non-python) solutions for mobile in place, e.g. tflite, libtorch, and onnxruntime.

> I wish there was more information as to how this differs from or improves on Jax.

The quick start guide has an overview.

> The Python API closely follows NumPy with a few exceptions. MLX also has a fully featured C++ API which closely follows the Python API.

The main differences between MLX and NumPy are:

        Composable function transformations: MLX has composable function transformations for automatic differentiation, automatic vectorization, and computation graph optimization.

        Lazy computation: Computations in MLX are lazy. Arrays are only materialized when needed.

        Multi-device: Operations can run on any of the supported devices (CPU, GPU, …)

The design of MLX is inspired by frameworks like PyTorch, Jax, and ArrayFire. A noteable difference from these frameworks and MLX is the unified memory model. Arrays in MLX live in shared memory. Operations on MLX arrays can be performed on any of the supported device types without performing data copies. Currently supported device types are the CPU and GPU.

https://ml-explore.github.io/mlx/build/html/index.html

Given that Apple has built PyTorch for Metal, I would assume they handle unified memory properly in there too… Same for Tensorflow and Jax on Metal.

I was excited about JAX, but I think the developers missed a trick when they decided it should be entirely immutable. It sounds silly, but I think if I have an array `x` and want to set index 0 to 10, it's a big mistake if I can't do:

  x[0] = 10

And instead I have to do:

  y = x.at[0].set(10)

Of course this has advantages, and I know it sounds lame, but as someone whose brain works in numpy, this is really offputting.

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.

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 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.

Of course `do` itself is much more capable, but it has an effect of the conversion for some monads, which was what I wanted to say.

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.

JAX is a wrapper on top of XLA. Instead of writing pure python, you're writing JAX abstractions.

for ex, a simple loop in JAX:

  def solve(i, v): return i+v
  x = jax.lax.fori_loop(0, 5, solve, 10)

Reply on news.ycombinator.com