Active Hackathon

Google open sources Switch Transformer models in T5X/JAX

The models include the 1.6T param Switch-C and the 395B param Switch-XXL.
Google open sources Switch Transformer models in T5X/JAX
Listen to this story

Google Brain has open sourced Switch Transformer models including 1.6T param Switch-C and the 395B param Switch-XXL in T5X/JAX.

Check out the GitHub repository here


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

What is JAX?

JAX (Just After eXecution) is a machine/deep learning library developed by DeepMind. All JAX operations are based on XLA or Accelerated Linear Algebra. XLA, developed by Google, is a domain-specific compiler for linear algebra that uses whole-program optimisations to accelerate computing. XLA makes BERT’s training speed faster by almost 7.3 times.


T5X is a modular, composable, research-friendly framework for high-performance, configurable, and inference of sequence models (starting with language) at many scales.

T5X allows many different classes of language tasks to be expressed uniformly, and a single encoder-decoder architecture can handle them without any task-specific parameters. 

Switch Transformer 

In deep learning, models usually reuse the same parameters for all inputs. However, the Switch Transformer uses a Mixture of Experts (MoE) and selects different parameters for each incoming example.

The model is primarily used for NLP research. The Transformer uses an algorithm called Switch Routing. Instead of activating several experts and combining the output, it chooses a single expert to work on an input. The algorithm simplifies the routing computation and reduces communication costs since individual expert models are hosted on different GPU devices.

More Great AIM Stories

Tasmia Ansari
Tasmia is a tech journalist at AIM, looking to bring a fresh perspective to emerging technologies and trends in data science, analytics, and artificial intelligence.

Our Upcoming Events

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

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

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

3 Ways to Join our Community

Discord Server

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

Telegram Channel

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

Subscribe to our newsletter

Get the latest updates from AIM