From 53dc786bc483042c9bfd9c6a0de4496e5a87302d Mon Sep 17 00:00:00 2001 From: Jake VanderPlas Date: Mon, 30 Jun 2025 13:28:22 -0700 Subject: [PATCH 1/2] pin tensorflow to nightly version --- .github/workflows/nightly.yaml | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/nightly.yaml b/.github/workflows/nightly.yaml index a3c5389..64758c0 100644 --- a/.github/workflows/nightly.yaml +++ b/.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 diff --git a/pyproject.toml b/pyproject.toml index 26d13f4..bbf85a9 100644 --- a/pyproject.toml +++ b/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", ] From e5bf9eb2d1ba4861b9e83eebaf97c1da176576a4 Mon Sep 17 00:00:00 2001 From: Jake VanderPlas Date: Mon, 30 Jun 2025 13:38:37 -0700 Subject: [PATCH 2/2] Update other package versions --- pyproject.toml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index bbf85a9..cf1eaf6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,11 +16,11 @@ keywords = [] # pip dependencies of the project dependencies = [ "chex==0.1.89", - "grain==0.2.9", - "jax==0.6.0", + "grain==0.2.10", + "jax==0.6.2", "flax==0.10.6", "ml_dtypes==0.5.1", - "optax==0.2.4", + "optax==0.2.5", "orbax-checkpoint==0.11.13", "orbax-export==0.0.6", ] @@ -41,7 +41,7 @@ dev = [ # TensorFlow datasets is an extra because it has a large install footprint. tfds = [ "tf-nightly", - "tensorflow_datasets==4.9.8", + "tensorflow_datasets==4.9.9", ] # Grain is now part of the default installation; keep this for