Merge pull request #5081 from google:fix-main-2

PiperOrigin-RevId: 830558189

Author: Flax Authors

.github/workflows/flax_test.yml

@@ -76,6 +76,8 @@ jobs:
     - name: Install standalone dependencies only
       run: |
         uv sync
+        # temporary: install jax nightly
+        uv pip install -U --pre jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
     - name: Test importing Flax
       run: |
         uv run python -c "import flax"
@@ -109,6 +111,7 @@ jobs:
 
     - name: Install dependencies
       run: |
+        rm -fr .venv
         uv sync --extra testing --extra docs
     - name: Install JAX
       run: |
@@ -122,13 +125,21 @@ jobs:
         if [[ "${{ matrix.test-type }}" == "doctest" ]]; then
           # TODO(cgarciae): Remove this once dm-haiku 0.0.14 is released
           uv pip install -U git+https://github.com/google-deepmind/dm-haiku.git
+          # temporary: install jax nightly
+          uv pip install -U --pre jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
           uv run tests/run_all_tests.sh --only-doctest
         elif [[ "${{ matrix.test-type }}" == "pytest" ]]; then
           uv pip install -U tensorflow-datasets
+          # temporary: install jax nightly
+          uv pip install -U --pre jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
           uv run tests/run_all_tests.sh --only-pytest
         elif [[ "${{ matrix.test-type }}" == "pytype" ]]; then
+        # temporary: install jax nightly
+          uv pip install -U --pre jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
           uv run tests/run_all_tests.sh --only-pytype
         elif [[ "${{ matrix.test-type }}" == "mypy" ]]; then
+          # temporary: install jax nightly
+          uv pip install -U --pre jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
           uv run tests/run_all_tests.sh --only-mypy
         else
           echo "Unknown test type: ${{ matrix.test-type }}"

.readthedocs.yml

@@ -1,31 +1 @@
-# .readthedocs.yml
-# Read the Docs configuration file
-# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
-
-# Required
-version: 2
-
-build:
-  os: ubuntu-22.04
-  tools:
-    python: "3.11"
-
-# Build documentation in the docs_nnx/ directory with Sphinx
-sphinx:
-  configuration: docs_nnx/conf.py
-
-# Optionally build your docs in additional formats such as PDF and ePub
-formats:
-  - htmlzip
-  - epub
-  # - pdf
-
-# Optionally set the version of Python and requirements required to build your docs
-python:
-  install:
-    - method: pip
-      path: .
-      extra_requirements:
-        - all
-        - testing
-        - docs
+# deprecated

docs/.readthedocs.yaml

@@ -8,7 +8,7 @@ version: 2
 build:
   os: ubuntu-22.04
   tools:
-    python: "3.11"
+    python: "3.12"
 
 # Build documentation in the docs/ directory with Sphinx
 sphinx:

docs_nnx/.readthedocs.yaml

@@ -8,7 +8,11 @@ version: 2
 build:
   os: ubuntu-22.04
   tools:
-    python: "3.11"
+    python: "3.12"
+  jobs:
+    pre_build:
+      - pip install ".[all, testing, docs]"
+      - pip install -U --pre jax jaxlib -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/
 
 # Build documentation in the docs/ directory with Sphinx
 sphinx:
@@ -19,13 +23,3 @@ formats:
   - htmlzip
   - epub
   # - pdf
-
-# Optionally set the version of Python and requirements required to build your docs
-python:
-  install:
-    - method: pip
-      path: .
-      extra_requirements:
-        - all
-        - testing
-        - docs