'Feedback from the JAX community'

[patrick-kidger]

Hi there! I found this project via the announcement in jax-ml/jax#28762. I'm following up on the offer of feedback from the JAX community. I'm not sure where you'd most prefer that to happen, feel free to redirect me from the issues page if you have a different venue in mind.

I had a few questions looking at the project!

One, is there any documentation on the PyTorch-mode vs JAX-mode? In particular how they might interact in a single unified way. I've now had quite a bit of experience in both frameworks individually, and I can't really imagine ever wanting the PyTorch mode. (Even though it's PyTorch I use in my day job 😁 !)

Second, is there a plan to ship a .nn library as part of the project? I think this has been the one big place where JAX went wrong -- the lack of a built-in library has sharded the community, where everyone has their own favourite alternative.
If so, I'm very curious whether you'll aim for a PyTorch-like OOP style (boo, hiss) or an Equinox-like functional style (of which I am unsurprisingly a fan)?

Third, I'd be curious to see a comparison between this and JAX? Partly the language, obviously, but at the framework level do you have a plan for how this will feel distinct? (If it helps then I have a laundry list of little mistakes made in JAX, that cannot now be changed.)

On to other nittier stuff: yikes, does Mojo not have variadic generics?? Even if it's just for tuple types?

[TilliFe]

Thanks for your valuable feedback!

• Re: Dual-mode execution (PyTorch-mode vs. JAX-mode):
  Nabla currently embraces a lazy mode of execution (no eager execution yet). We leverage Mojo's memory management—where values are destroyed upon last use—to track value lifetimes. This allows Nabla to determine if/when to materialize values and to include intermediates (like optimizer states) in the compiled graph by observing reference counts. In other words, when we 'need' a value (triggered by a load/print statement), we potentially compute more values than those that we actively load or print, allowing for more interesting programming patterns. I hope this make a bit more sense now.?
• Re: .nn module:
  Excellent suggestion. We're focused on building a robust core before adding an nn module, though I agree that nn modules should ideally be a part of sth that calls itself an ML library. Mojo's relatively small binaries might reduce "bloat" concerns, making this feasible. Regarding fragmentation: from your JAX experience, when does it become a hurdle, and at which point are specialized libraries appropriate?
• Re: Distinctive qualities compared to JAX:
  Nabla is built in Mojo for the Mojo ecosystem (no Python port planned at the moment, however this might change). Hopefully, Mojo's Pythonic evolution will lower the entry barrier for Python users in the future. Moreover, a key differentiator in my mind will be Mojo's built-in capabilities for low-level kernel-writing for fine-grained control and optimization. Modular promotes Mojo as being hardware-agnostic; and I'm fairly optimistic about this. Your "JAX mistakes" list is intriguing! I'd love to hear those insights.
  It should be mentioned, that Nabla is clearly in an early stage, let's see where the journey goes.
• Re: Variadic generics:
  Good point. This could be better! 😆

Thanks again. Please share any further thoughts!

[patrick-kidger]

Great stuff!

• On dual mode execution: hmm, not entirely. I guess it's not clear to me how loss = some_fn(...); loss.backward() interacts with JAX-style loss = jax.grad(some_fn)(...). What happens if the two approaches are mix-and-matched? What is the underlying abstraction that a user should have in their head?

• On nn: makes sense to hold off on this until more pieces in place! As for scope, I think the PyTorch .nn, .optim, .utils.data namespaces are probably the optimal point for what to include built-in (vs what to leave to other libraries). This ensures that the framework comes 'batteries included' for the 95% ML use-case. (Implementation-wise note that I think Optax does a better job than torch.optim and Grain offers a better API than torch.utils.data.)

• On comparing to JAX: if kernel-writing is the main point of comparison then consider me as updating my question to 'vs JAX+Pallas?' (Notably the latter does a much better job than Triton due to the way it supports jax.vmap etc transformations, and for parameterisastion doesn't require codegen.)

• As for JAX mistakes! Here's a top few that come to mind:

  • jax.jit and friends will cast Python scalars to scalar-shaped arrays, the latter of which are then traced. This introduces a whole host of downstream thorniness: for example this is the reason that 'weak dtypes' are needed, so that my_array + my_scalar will perform the same dtype promotion regardless of whether or not my_scalar has been traced, i.e. whether my_array.dtype is used directly or whether a join over the dtype lattice is performed.
  • Not properly supporting runtime errors, either as an abort or as an algebraic effect.
  • No way to zero-copy flatten/unflatten a pytree. For example if one has a tree-of-arrays and would like to call a QR decomposition that expects specifically a flat array. Ideally the arrays were already in contiguous memory and no copy would be needed to concatenate them.
  • Not including a dataclass+pytree combo built-in to the library. Everyone ends up reinventing their own version of this with its many footguns typically not well-handled in a central place.
  • Bound methods and functools.partial should both be pytrees. The former is a footgun, the latter leads to a proliferation of jax.tree_util.Partial, equinox.Partial etc to handle this limitation. (If you really wanted to push the boat out here then you could also arange for closures to be pytrees, and thus avoid the need for closure-conversion around higher-order primitives and custom autodiff rules.)

  The last few refer to 'pytrees'; I realise given nabla's possibility of integration with its host language that this might actually be a notion shared with the language rather than a new concept on top.

[TilliFe]

This is so great Patrick, thank you very much!

• Your skepticism about the dual nature of the library is appreciated! 👍 I agree that there could be unintended ambiguities between the two approaches. To be honest, I might eventually decide to remove the PyTorch-like typing style, i.e. the ability to explicitly change the computation graph. Personally, I also prefer the (purely) functional JAX-like approach.
  Nonetheless, I wonder how the PyTorch people deal with this ambiguity, since they also promote the use of their new functional module (formerly known as funcTorch). So I'll think about it. That kind of flexibility is probably a feature, not a bug, in the beginning.
• re., JAX + Pallas: I need to learn more about Pallas, it looks cool. However, Python - regardless of the DSL built on top of it - can't replace real bare-metal programming. Mojo is a systems programming language designed specifically for such purposes, and I think people will find it much more convenient to write low-level kernels in Mojo than if they do it in Python. Not sure, just a feeling ;). You are right to argue though that Pallas' deep integration into JAX' transformation system, sets the bar really really high!