Jacked About Jax

As others have discussed, we’ve noticed a recent uptick in research implemented using Jax. You could chalk it up as yet another ML platform, but Jax is emerging as the tool of choice for faster research iteration at DeepMind. After exploring for ourselves, we’re excited to find Jax is principally designed for fast differentiation. Excited because differentiation is foundational to gradient-based learning strategies supported in many ML algorithms. Moreover, the derivative is also ubiquitous in scientific computing, making Jax one powerful hammer!...

 · 4 min · Terry Rodriguez