# JAX

AXLearn is a deep learning library built by Apple based on JAX and XLA. It takes an object-oriented approach to address software engineering challenges in large-scale deep learning model development. Its configuration system allows users to combine models from reusable building blocks and integrate with other libraries (such as Flax and Hugging Face transformers). AXLearn aims to scale training, supporting models with hundreds of billions of parameters efficiently trained on thousands of accelerators, suitable for deployment on public clouds. It also adopts a global computation paradigm, allowing users to describe computations on a global virtual machine instead of individual accelerators. AXLearn supports a wide range of applications, including natural language processing, computer vision, and speech recognition, and includes baseline configurations for training state-of-the-art models.

JaxMARL is a multi-agent reinforcement learning library that combines ease of use with GPU accelerated performance. It supports common multi-agent reinforcement learning environments and popular benchmark algorithms. The goal is to provide a comprehensive library for evaluating multi-agent reinforcement learning methods and comparing them with related benchmarks. It also introduces SMAX, a simplified version of the popular StarCraft Multi-Agent Challenge environment, which does not require running the StarCraft II game engine.