From 218f9f694848e6bb964c7f7c317d2348d64d4e65 Mon Sep 17 00:00:00 2001 From: Toshi Pahadia Date: Mon, 5 Oct 2026 11:25:26 +0530 Subject: [PATCH] Raise flax floor to 0.12.10 so nightly images import under JAX >= 0.11 setup.sh installs generated_requirements.txt with `uv pip install --resolution=lowest`, so whatever floor is written there is exactly what ships in the image. The floor was `flax>=0.12.6`; that release still subclasses `jax.core.Effect` (flax/nnx/variablelib.py:289), which JAX removed in 0.11.0. The nightly step then upgrades JAX to 0.11.x, and the result cannot import flax.nnx at all: AttributeError: jax.core.Effect was deprecated in JAX v0.10.0 and removed in JAX v0.11.0. Use jax.extend.core.Effect. Every maxdiffusion entry point imports flax.nnx, so the official nightly image built from main (maxdiffusion_jax_nightly:2026-10-03 .. :2026-10-05, flax 0.12.6 + jax 0.11.2) fails at startup for BF16 and FP8 alike. Images built from a tree where the floor had been raised resolved flax 0.12.10 and passed the Ironwood 4x4x4 Wan 2.1 runs, which is why the breakage looked intermittent: it depends on which floor the image was built from, not on the day. flax 0.12.7 is the first release without jax.core.Effect; 0.12.9+ declare jax>=0.11.1. 0.12.10 is the version verified end-to-end on hardware, and with --resolution=lowest the floor is the shipped version, so pin the floor to the tested one. Hand-edited (generated + base requirements, deps table) to avoid seed-env churn, same as #499. b/537854580 --- dependencies/requirements/base_requirements/requirements.txt | 2 +- .../requirements/generated_requirements/requirements.txt | 2 +- src/maxdiffusion/dependency_versions_table.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/dependencies/requirements/base_requirements/requirements.txt b/dependencies/requirements/base_requirements/requirements.txt index b2a5d239c..a446ae034 100644 --- a/dependencies/requirements/base_requirements/requirements.txt +++ b/dependencies/requirements/base_requirements/requirements.txt @@ -6,7 +6,7 @@ av chex datasets einops -flax +flax>=0.12.10 ftfy google-cloud-storage grain diff --git a/dependencies/requirements/generated_requirements/requirements.txt b/dependencies/requirements/generated_requirements/requirements.txt index da356c936..0e66b4637 100644 --- a/dependencies/requirements/generated_requirements/requirements.txt +++ b/dependencies/requirements/generated_requirements/requirements.txt @@ -41,7 +41,7 @@ etils>=1.13.0 execnet>=2.1.2 filelock>=3.20.3 flatbuffers>=25.12.19 -flax>=0.12.6 +flax>=0.12.10 fonttools>=4.61.1 frozenlist>=1.8.0 fsspec>=2026.1.0 diff --git a/src/maxdiffusion/dependency_versions_table.py b/src/maxdiffusion/dependency_versions_table.py index 241035ac3..b6074e551 100644 --- a/src/maxdiffusion/dependency_versions_table.py +++ b/src/maxdiffusion/dependency_versions_table.py @@ -8,7 +8,7 @@ "aqtp": "aqtp", "datasets": "datasets", "einops": "einops", - "flax": "flax", + "flax": "flax>=0.12.10", "ftfy": "ftfy", "google-cloud-storage": "google-cloud-storage", "grain": "grain",