What Are DeepMind’s Newly Released Libraries For Neural Networks & Reinforcement Learning?

AI research startup DeepMind has now open-sourced new libraries for neural networks and reinforcement learning based on JAX. JAX is a numerical computing library launched by Google a couple of years ago, and can automatically differentiate native Python and NumPy functions. JAX uses XLA (Accelerated Linear Algebra) to compile and run your NumPy programs on GPUs and TPUs, which is great for machine learning to compute. 

Known as Haiku and RLax, both open-sourced libraries can be used by AI enthusiasts, professionals and researchers for reinforcement learning initiatives.

Haiku: Sonnet for JAX 

This Haiku library allows AI researchers to fully access JAX’s pure function transformations by providing object-oriented programming models. Haiku is a neural network library for JAX, meaning it is created on the programming model and APIs of Sonnet– a library on top of TensorFlow designed to provide simple, composable abstractions for machine learning research. Haiku, in fact, has been developed by some of the authors of Sonnet. The open-source library Sonnet is widely adopted at DeepMind.


Sign up for your weekly dose of what's up in emerging technology.

Haiku has been experimented by researchers at DeepMind at scale with success. Haiku has demonstrated promising results in big-scale experimentation on image and language datasets using reinforcement learning. But, Haiku isn’t the only neural network library made for JAX, and others also exist such as Google’ Flax.

RLax, on the other hand, is a simple RL library for JAX. Instead of providing complete algorithms for RL, this library provides useful building blocks for implementing specific mathematical operations to build fully-functional RL agents. DeepMind pioneered deep reinforcement learning – to create the first artificial agents to achieve human-level performance across many challenging domains. Two years ago, the company rolled out the first widely successful algorithm for deep reinforcement learning. 

Advantages Of New DeepMind Open-Source Libraries

So, if there are many neural network libraries for JAX, why should you choose Haiku and RLax?

According to Deep Mind, Haiku has been created to make specific processes straightforward for managing things like model parameters along with the model state.

Haiku is presently in alpha, many researchers have experimented with Haiku for many months and have reproduced a number of experiments at a large scale.

DeepMind has already reproduced many experiments in Haiku and JAX, and include big-scale tests in image and language processing models, generative models, as well as reinforcement learning.

The advantage is that both libraries have well-defined instructions, and for Haiku, the APIs are similar to Sonnet, which means moving to Haiku is likely to be quite easy for those working on Sonnet. The fact that Haiku and RLax are supported and tested by Deep Mind can give developers, researchers and students peace of mind for their advanced ML endeavours.


It seems while project Haiku and RLax have been launched to GitHub, they are still experimental. While JAX had been created for high-performance machine learning research by giving an integrated system for learning model optimisation, it is not the most comfortable tool to operate. This is why experts think DeepMind developed Haiku and RLax.

Also Read: 10 Best ML Engineering Practices

More Great AIM Stories

Vishal Chawla
Vishal Chawla is a senior tech journalist at Analytics India Magazine and writes about AI, data analytics, cybersecurity, cloud computing, and blockchain. Vishal also hosts AIM's video podcast called Simulated Reality- featuring tech leaders, AI experts, and innovative startups of India.

Our Upcoming Events

Conference, in-person (Bangalore)
Machine Learning Developers Summit (MLDS) 2023
19-20th Jan, 2023

Conference, in-person (Bangalore)
Rising 2023 | Women in Tech Conference
16-17th Mar, 2023

Conference, in-person (Bangalore)
Data Engineering Summit (DES) 2023
27-28th Apr, 2023

Conference, in-person (Bangalore)
MachineCon 2023
23rd Jun, 2023

Conference, in-person (Bangalore)
Cypher 2023
20-22nd Sep, 2023

3 Ways to Join our Community

Whatsapp group

Discover special offers, top stories, upcoming events, and more.

Discord Server

Stay Connected with a larger ecosystem of data science and ML Professionals

Subscribe to our newsletter

Get the latest updates from AIM