Merge branch 'jax-ml:main' into upgrade-vit-2

Author: 8bitmp3

.github/workflows/lint.yaml

@@ -23,7 +23,7 @@ jobs:
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
       - name: Set up Python 3.12
-        uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55  # v5.5.0
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065  # v5.6.0
         with:
           python-version: 3.12
       - uses: pre-commit/action@2c7b3805fd2a0fd8c1884dcaebf91fc102a13ecd  # v3.0.1

.github/workflows/nightly.yaml

@@ -31,7 +31,7 @@ jobs:
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55  # v5.5.0
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065  # v5.6.0
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
@@ -56,13 +56,13 @@ jobs:
       fail-fast: false
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
-      - uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55  # v5.5.0
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065  # v5.6.0
         with:
           python-version: 3.12
       - name: Install dependencies with jax nightly
         run: |
           python -m pip install --upgrade pip
-          python -m pip install .[dev,tfds,grain]
+          python -m pip install .[dev,tfds]
           python -m pip install --upgrade --pre jax jaxlib -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
       - name: Run tests
         run: |
@@ -81,13 +81,13 @@ jobs:
       fail-fast: false
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
-      - uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55  # v5.5.0
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065  # v5.6.0
         with:
           python-version: 3.12
       - name: Install dependencies with flax nightly
         run: |
           python -m pip install --upgrade pip
-          python -m pip install .[dev,tfds,grain]
+          python -m pip install .[dev,tfds]
           python -m pip install --upgrade git+https://github.com/google/flax.git
       - name: Run tests
         run: |
@@ -106,13 +106,13 @@ jobs:
       fail-fast: false
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
-      - uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55  # v5.5.0
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065  # v5.6.0
         with:
           python-version: 3.12
       - name: Install dependencies with optax nightly
         run: |
           python -m pip install --upgrade pip
-          python -m pip install .[dev,tfds,grain]
+          python -m pip install .[dev,tfds]
           python -m pip install --upgrade git+https://github.com/google-deepmind/optax.git
       - name: Run tests
         run: |
@@ -131,13 +131,13 @@ jobs:
       fail-fast: false
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
-      - uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55  # v5.5.0
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065  # v5.6.0
         with:
           python-version: 3.12
       - name: Install dependencies with orbax-checkpoint and orbax-export nightly
         run: |
           python -m pip install --upgrade pip
-          python -m pip install .[dev,tfds,grain]
+          python -m pip install .[dev,tfds]
           python -m pip install --upgrade 'git+https://github.com/google/orbax/#subdirectory=checkpoint' 'git+https://github.com/google/orbax/#subdirectory=export'
       - name: Run tests
         run: |
@@ -156,13 +156,13 @@ jobs:
       fail-fast: false
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
-      - uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55  # v5.5.0
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065  # v5.6.0
         with:
           python-version: 3.12
       - name: Install dependencies with chex nightly
         run: |
           python -m pip install --upgrade pip
-          python -m pip install .[dev,tfds,grain]
+          python -m pip install .[dev,tfds]
           python -m pip install --upgrade 'git+https://github.com/google-deepmind/chex/'
       - name: Run tests
         run: |

.github/workflows/release.yml

@@ -43,7 +43,7 @@ jobs:
     permissions:
       id-token: write
     steps:
-      - uses: actions/download-artifact@95815c38cf2ff2164869cbab79da8d1f422bc89e  # v4.2.1
+      - uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093  # v4.3.0
         with:
           name: distribution
           path: dist

.github/workflows/test.yaml

@@ -31,13 +31,13 @@ jobs:
     steps:
     - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@8d9ed9ac5c53483de85588cdf95a591a75ab9f55  # v5.5.0
+      uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065  # v5.6.0
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install .[dev,tfds,grain]
+        pip install .[dev,tfds]
     - name: Run tests
       run: |
         pytest -n auto jax_ai_stack

README.md

@@ -23,8 +23,7 @@ single point-of-entry for this suite of libraries, so you can install and begin
 using many of the same open source packages that Google developers are using
 in their everyday work.
 
-To get started with the JAX AI stack, you can check out [Getting started with JAX](
-https://github.com/jax-ml/jax-ai-stack/blob/main/docs/source/getting_started_with_jax_for_AI.ipynb).
+To get started with the JAX AI stack, you can check out [Getting started with JAX](https://docs.jaxstack.ai/en/latest/getting_started.html).
 This is still a work-in-progress, please check back for more documentation and tutorials
 in the coming weeks!
 
@@ -44,20 +43,27 @@ together via the integration tests in this repository. Packages include:
 - [optax](https://github.com/google-deepmind/optax): gradient processing and optimization in JAX.
 - [orbax](https://github.com/google/orbax): checkpointing and persistence utilities for JAX.
 - [chex](https://github.com/google-deepmind/chex): utilities for writing reliable JAX code.
+- [grain](https://github.com/google/grain): data loading.
 
 ### Optional packages
 
 Additionally, there are optional packages you can install with `pip` extras.
+
 The following command:
 ```
-pip install jax-ai-stack[grain]
+pip install jax-ai-stack[tfds]
 ```
-will install a compatible version of the [grain](https://github.com/google/grain) data
-loader (currently mac and linux-only).
+will install a compatible version of
+[tensorflow](https://github.com/tensorflow/tensorflow)
+and [tensorflow-datasets](https://github.com/tensorflow/datasets).
 
-Similarly, the following command:
+### Hardware support
+
+To install `jax-ai-stack` with hardware-specific JAX support, add the JAX installation
+command in the same `pip install` invocation. For example:
 ```
-pip install jax-ai-stack[tfds]
+pip install jax-ai-stack "jax[cuda]"  # JAX + AI stack with GPU/CUDA support
+pip install jax-ai-stack "jax[tpu]"  # JAX + AI stack with TPU support
 ```
-will install a compatible version of [tensorflow](https://github.com/tensorflow/tensorflow)
-and [tensorflow-datasets](https://github.com/tensorflow/datasets).
+For more information on available options for hardware-specific JAX installation, refer
+to [JAX installation](https://docs.jax.dev/en/latest/installation.html).