Metadata-Version: 2.1
Name: axlearn
Version: 0.0.1
Summary: AXLearn
Requires-Python: >=3.8
Description-Content-Type: text/markdown
Requires-Dist: attrs>=21.3.0
Requires-Dist: absl-py
Requires-Dist: chex>=0.1.6
Requires-Dist: flax>=0.7.1
Requires-Dist: importlab==0.7
Requires-Dist: jax==0.4.13
Requires-Dist: jaxlib==0.4.13
Requires-Dist: nltk==3.7
Requires-Dist: numpy<1.24
Requires-Dist: optax==0.1.3
Requires-Dist: portpicker
Requires-Dist: protobuf==3.20.3
Requires-Dist: tensorboard-plugin-profile==2.8.0
Requires-Dist: tensorflow==2.8.0; platform_machine == 'x86_64'
Requires-Dist: tensorflow-datasets==4.7.0
Requires-Dist: tensorflow-io[tensorflow]; platform_machine == 'x86_64'
Requires-Dist: tensorflow-io-gcs-filesystem[tensorflow]; platform_machine == 'x86_64'
Requires-Dist: tensorflow_text==2.8.2
Requires-Dist: tensorstore>=0.1.21
Requires-Dist: toml
Requires-Dist: typing-extensions>=4.2.0
Requires-Dist: seqio==0.0.15
Requires-Dist: attrs>=21.3.0 ; extra == "apple-silicon"
Requires-Dist: absl-py ; extra == "apple-silicon"
Requires-Dist: chex>=0.1.6 ; extra == "apple-silicon"
Requires-Dist: flax>=0.7.1 ; extra == "apple-silicon"
Requires-Dist: jax==0.4.13 ; extra == "apple-silicon"
Requires-Dist: jaxlib==0.4.13 ; extra == "apple-silicon"
Requires-Dist: nltk==3.7 ; extra == "apple-silicon"
Requires-Dist: optax>=0.1.1 ; extra == "apple-silicon"
Requires-Dist: portpicker ; extra == "apple-silicon"
Requires-Dist: tensorboard-plugin-profile==2.8.0 ; extra == "apple-silicon"
Requires-Dist: tensorflow-datasets==4.7.0 ; extra == "apple-silicon"
Requires-Dist: tensorstore>=0.1.21 ; extra == "apple-silicon"
Requires-Dist: typing-extensions>=4.2.0 ; extra == "apple-silicon"
Requires-Dist: seqio==0.0.15 ; extra == "apple-silicon"
Requires-Dist: black==23.1a1 ; extra == "dev"
Requires-Dist: diffusers==0.16.1 ; extra == "dev"
Requires-Dist: einops ; extra == "dev"
Requires-Dist: evaluate ; extra == "dev"
Requires-Dist: fairseq==0.12.2 ; extra == "dev"
Requires-Dist: isort ; extra == "dev"
Requires-Dist: pre-commit ; extra == "dev"
Requires-Dist: pycocotools ; extra == "dev"
Requires-Dist: pylint==2.13.9 ; extra == "dev"
Requires-Dist: pytest ; extra == "dev"
Requires-Dist: pytest-xdist ; extra == "dev"
Requires-Dist: pytype==2022.4.22 ; extra == "dev"
Requires-Dist: scikit-learn>=1.1.1 ; extra == "dev"
Requires-Dist: scipy ; extra == "dev"
Requires-Dist: sentencepiece != 0.1.92 ; extra == "dev"
Requires-Dist: tqdm ; extra == "dev"
Requires-Dist: timm==0.6.12 ; extra == "dev"
Requires-Dist: torch>=1.12.1 ; extra == "dev"
Requires-Dist: torchvision==0.14.1 ; extra == "dev"
Requires-Dist: transformers==4.27.1 ; extra == "dev"
Requires-Dist: wandb ; extra == "dev"
Requires-Dist: wrapt ; extra == "dev"
Requires-Dist: cloud-tpu-client ; extra == "gcp"
Requires-Dist: crcmod ; extra == "gcp"
Requires-Dist: google-api-python-client==1.8.0 ; extra == "gcp"
Requires-Dist: google-auth==1.35.0 ; extra == "gcp"
Requires-Dist: google-cloud-storage==2.2.1 ; extra == "gcp"
Requires-Dist: google-cloud-core==1.7.3 ; extra == "gcp"
Requires-Dist: jax[tpu]==0.4.13 ; extra == "tpu"
Requires-Dist: setuptools==65.7.0 ; extra == "vertexai_tensorboard"
Requires-Dist: google-cloud-aiplatform[tensorboard] ; extra == "vertexai_tensorboard"
Provides-Extra: apple-silicon
Provides-Extra: dev
Provides-Extra: gcp
Provides-Extra: tpu
Provides-Extra: vertexai_tensorboard

# The AXLearn Library for Deep Learning

AXLearn is a library built on top of [JAX](https://jax.readthedocs.io/) and
[XLA](https://www.tensorflow.org/xla) to support development of large-scale deep learning models.

AXLearn takes an object-oriented approach to the software engineering challenges that arise from
building, iterating, and maintaining models.
The configuration system of the library lets users compose models from reusable building blocks and
integrate with other libraries such as [Flax](https://flax.readthedocs.io/) and
[Hugging Face transformers](https://github.com/huggingface/transformers).

AXLearn is built to scale.
It supports training of models with up to hundreds of billions of parameters across thousands of
accelerators at high utilization.
It is also designed to run on public clouds and provides tools to deploy and manage jobs and data.
Built on top of [GSPMD](https://arxiv.org/abs/2105.04663), AXLearn adopts a global computation
paradigm to allow users to describe computation on a virtual global computer rather than on a
per-accelerator basis.

AXLearn supports a wide range of applications, including natural language processing, computer
vision, and speech recognition and contains baseline configurations for training state-of-the-art
models.

