Merge branch 'main' into sharding_metadata

Author: samanklesaria

.github/workflows/flax_test.yml

@@ -80,7 +80,7 @@ jobs:
         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"
+        uv run --no-sync python -c "import flax"
 
   tests:
     name: Run Tests
@@ -127,20 +127,20 @@ jobs:
           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
+          uv run --no-sync 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
+          uv run --no-sync tests/run_all_tests.sh --only-pytest
         elif [[ "${{ matrix.test-type }}" == "pytype" ]]; then
-        # temporary: install jax nightly
+          # 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
+          uv run --no-sync 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
+          uv run --no-sync tests/run_all_tests.sh --only-mypy
         else
           echo "Unknown test type: ${{ matrix.test-type }}"
           exit 1

docs_nnx/api_reference/flax.nnx/graph.rst

@@ -31,7 +31,11 @@ graph
 
 .. autofunction:: find_duplicates
 .. autofunction:: pure
-.. autofunction:: to_refs
-.. autofunction:: to_arrays
+.. autofunction:: as_immutable_vars
+.. autofunction:: as_mutable_vars
+.. autofunction:: as_hijax_vars
+.. autofunction:: as_pytree_vars
+.. autofunction:: as_ref_vars
+.. autofunction:: as_array_vars
 .. autofunction:: flatten
 .. autofunction:: unflatten