3 ms·

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

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.

leewlving 3y ago
I think it is a feature from functional programming. In some level, I more agree with jax style.
- lifthrasiir 3y ago
  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.
  - n2d4 3y ago
    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.
    - lifthrasiir 3y ago
      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.
boywitharupee 3y ago
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)

leewlving 3y ago

I think it is a feature from functional programming. In some level, I more agree with jax style.

lifthrasiir 3y ago
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.
- n2d4 3y ago
  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.
  - lifthrasiir 3y ago
    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.

boywitharupee 3y ago

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)