Merge pull request #4909 from google:fix-mypy

PiperOrigin-RevId: 799720766

Author: Flax Authors

.github/workflows/flax_test.yml

@@ -72,7 +72,7 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - uses: astral-sh/setup-uv@887a942a15af3a7626099df99e897a18d9e5ab3a # v5.1.0
       with:
-        uv-version: "0.3.0"
+        version: "0.8.13"
     - name: Install standalone dependencies only
       run: |
         uv sync
@@ -108,7 +108,7 @@ jobs:
     - name: Setup uv
       uses: astral-sh/setup-uv@887a942a15af3a7626099df99e897a18d9e5ab3a # v5.1.0
       with:
-        version: "0.3.0"
+        version: "0.8.13"
 
     - name: Install dependencies
       run: |

.github/workflows/jax_nightly.yml

@@ -34,7 +34,7 @@ jobs:
     - name: Setup uv
       uses: astral-sh/setup-uv@887a942a15af3a7626099df99e897a18d9e5ab3a  # v5.1.0
       with:
-        version: "0.3.0"
+        version: "0.8.13"
     - name: Install dependencies
       run: |
         uv sync --extra testing --extra docs

flax/nnx/graph.py

@@ -73,7 +73,7 @@ class Repeated: ...
 @jax.tree_util.register_dataclass
 @dataclasses.dataclass(frozen=True, slots=True, repr=False)
 class ArrayRefOutput(reprlib.Representable):
-  value: jax.Array | NoUpdate | Repeated
+  value: jax.Array
 
   def __nnx_repr__(self):
     yield reprlib.Object(type=type(self))

flax/nnx/rnglib.py

@@ -840,7 +840,7 @@ def split_rngs_wrapper(*args, **kwargs):
       if squeeze:
         key = key[0]
       if variablelib.is_array_ref(stream.key.raw_value):
-        stream.key.raw_value = variablelib.array_ref(key)
+        stream.key.raw_value = variablelib.array_ref(key)  # type: ignore[assignment]
       else:
         stream.key.value = key
       if squeeze:
@@ -852,7 +852,7 @@ def split_rngs_wrapper(*args, **kwargs):
 
       count = jnp.zeros(counts_shape, dtype=jnp.uint32)
       if variablelib.is_array_ref(stream.count.raw_value):
-        stream.count.raw_value = variablelib.array_ref(count)
+        stream.count.raw_value = variablelib.array_ref(count)  # type: ignore[assignment]
       else:
         stream.count.value = count
 

flax/nnx/variablelib.py

@@ -247,7 +247,7 @@ def __init__(
       if is_array_ref(value):
         _value = tp.cast(A, value)
       else:
-        _value = array_ref(jnp.asarray(value))
+        _value = array_ref(jnp.asarray(value))  # type: ignore[assignment]
     else:
       _value = value