ci: unblock Flax tests during JAX compatibility window (#10646)

Author: Light2Dark

pyproject.toml

@@ -221,7 +221,8 @@ test-optional = [
     "jupytext>=1.17.2",
     # For standard scientific computing/ ML
     # flax 0.10 requires jax>=0.4.27, so keep the floors compatible.
-    "jax>=0.4.27; python_version == '3.12'",
+    # TODO: Remove <0.11.1 after Flax 0.12.9 enters CI's 7-day dependency window.
+    "jax>=0.4.27,<0.11.1; python_version == '3.12'",
     "flax>=0.10.0; python_version == '3.12'",
     "torch>=2.4.0; python_version == '3.12'",
     "scikit-bio>=0.6.3; python_version == '3.12'",