Make grain part of the default installation (#202)

Author: jakevdp

.github/workflows/nightly.yaml

@@ -62,7 +62,7 @@ jobs:
       - 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: |
@@ -87,7 +87,7 @@ jobs:
       - 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: |
@@ -112,7 +112,7 @@ jobs:
       - 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: |
@@ -137,7 +137,7 @@ jobs:
       - 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: |
@@ -162,7 +162,7 @@ jobs:
       - 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/test.yaml

@@ -37,7 +37,7 @@ jobs:
     - 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

@@ -44,20 +44,16 @@ 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]
-```
-will install a compatible version of the [grain](https://github.com/google/grain) data
-loader (currently mac and linux-only).
 
-Similarly, the following command:
+The following command:
 ```
 pip install jax-ai-stack[tfds]
 ```
-will install a compatible version of [tensorflow](https://github.com/tensorflow/tensorflow)
+will install a compatible version of
+[tensorflow](https://github.com/tensorflow/tensorflow)
 and [tensorflow-datasets](https://github.com/tensorflow/datasets).

docs/source/install.md

@@ -14,18 +14,12 @@ together via the integration tests in this repository. Packages include:
 - [ml_dtypes](http://github.com/jax-ml/ml_dtypes): NumPy dtype extensions for machine learning.
 - [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]
-```
-will install a compatible version of the [grain](https://github.com/google/grain) data
-loader (currently linux-only).
-
-Similarly, the following command:
+Additionally, there are optional packages you can install with `pip` extras. The following command:
 ```
 pip install jax-ai-stack[tfds]
 ```

pyproject.toml

@@ -16,6 +16,7 @@ keywords = []
 # pip dependencies of the project
 dependencies = [
     "chex==0.1.89",
+    "grain==0.2.8",
     "jax==0.6.0",
     "flax==0.10.6",
     "ml_dtypes==0.5.1",
@@ -43,10 +44,9 @@ tfds = [
     "tensorflow_datasets==4.9.8",
 ]
 
-# Grain is an extra because as of v0.2.0 it has no OSX wheels.
-grain = [
-    "grain==0.2.8",
-]
+# Grain is now part of the default installation; keep this for
+# backward compatibility.
+grain = []
 
 [tool.pyink]
 # Formatting configuration to follow Google style-guide