pin tensorflow to nightly version

Author: jakevdp

.github/workflows/nightly.yaml

@@ -38,7 +38,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           python -m pip install pytest pytest-xdist
-          python -m pip install -U chex jax flax grain ml_dtypes optax orbax-checkpoint orbax-export tensorflow tensorflow_datasets
+          python -m pip install -U chex jax flax grain ml_dtypes optax orbax-checkpoint orbax-export tf-nightly tensorflow_datasets
       - name: Run tests
         run: |
           pytest -n auto jax_ai_stack

pyproject.toml

@@ -40,7 +40,7 @@ dev = [
 
 # TensorFlow datasets is an extra because it has a large install footprint.
 tfds = [
-    "tensorflow==2.19.0",
+    "tf-nightly",
     "tensorflow_datasets==4.9.8",
 ]