I'm curious about the reason why you use JAX. about dm-haiku HOT 1 CLOSED

Hey! DeepMind is a big place, so the reasons different researchers use JAX will depend on their individual use cases; the key features for metalearning aren't necessarily the same as for generative models.

I started working on Haiku with Tom to unblock my own projects. For my personal work:

• As you mention, JAX is lightweight. I'm able to see through the framework.
• JAX maps well to XLA's computational model. XLA is the primary pathway to using TPUs, which are a huge enabler for my research.
• vmap and pmap are elegant ways of achieving parallelism.
• Fast eager execution is useful for debugging.
• XLA allows JAX to express computations as compositions of simple primitives, rather than a collection of fused ops. This enables two things:
  1. It's easy to write code that's fast and looks like the math, all in Python.
  2. Since everything is a composition of primitives, higher-order and mixed-mode autodiff is much easier.

Again, I want to emphasize that DeepMind has a large portfolio of research and that the researchers who do use JAX use it for a variety of reasons.

from dm-haiku.