GitHub repository: superckid/Mixture-of-Depths-Jax
Primary language: Python
Description: Jax module for the paper: "Mixture-of-Depths: Dynamically allocating compute in transformer-based language models"