I'm curious about the reason why you use JAX.

[chris-chris]

After you open-sourced haiku, rlax recently, I started to learn about JAX and I love these projects. However, whenever I share these open sources to deep learning communities and my friends, they don't seem to be convinced to use them. I believe you must have solid reasons for choosing JAX for the RL research framework, which I've tried to guess. So far, these are the advantages that I could think of.

• Pure function without side effects
• Performance enhancement using XLA
• Lightweight and NEAT
• Fully utilizing vectorization and parallelization
• Flexible architecture for a complicated distributed learning system

But I'm still not fully aware of the logic behind it. Would it be possible for you to briefly explain the motivation behind using JAX for your research?

[trevorcai]

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.