Added workaround by P.Hawkins

Author: vfdev-5

.github/workflows/flax_test.yml

@@ -191,6 +191,8 @@ jobs:
       run: |
         rm -fr .venv
         uv sync --extra testing --extra docs
+        # 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 with pytest
       run: |
         export XLA_FLAGS='--xla_force_host_platform_device_count=4'

flax/linen/module.py

@@ -19,6 +19,7 @@
 import enum
 import functools
 import inspect
+import sys
 import threading
 import typing
 import weakref
@@ -1097,6 +1098,12 @@ def _customized_dataclass_transform(cls, kw_only: bool):
     ) in extra_fields:
       setattr(cls, name, default)
       cls.__annotations__[name] = annotation
+
+    # TODO: a workaround for the issue:
+    # https://github.com/google/flax/pull/5087#issuecomment-3536610568
+    if (sys.version_info.major, sys.version_info.minor) in [(3, 12), (3, 13)]:
+      setattr(cls, '__annotations__', cls.__annotations__)
+
     dataclasses.dataclass(  # type: ignore[call-overload]
       unsafe_hash='__hash__' not in cls.__dict__,
       repr=False,

flax/nnx/variablelib.py

@@ -1712,8 +1712,8 @@ def __contains__(self, item) -> bool:
 
   def __eq__(self, other) -> bool:
     if isinstance(other, Variable):
-      other = other[...]
-    return self[...].__eq__(other)  # type: ignore
+      other = other.get_value()
+    return self.get_value().__eq__(other)  # type: ignore
 
   def __iadd__(self: V, other) -> V:
     raise NotImplementedError(
@@ -1736,61 +1736,61 @@ def __imul__(self: V, other) -> V:
   def __imatmul__(self: V, other) -> V:
     raise NotImplementedError(
       'In-place operations are no longer supported for Variable.\n'
-      'Use `variable.value @= x` instead.'
+      'Use `variable[...] @= x` instead.'
     )
 
   def __itruediv__(self: V, other) -> V:
     raise NotImplementedError(
       'In-place operations are no longer supported for Variable.\n'
-      'Use `variable.value /= x` instead.'
+      'Use `variable[...] /= x` instead.'
     )
 
   def __ifloordiv__(self: V, other) -> V:
     raise NotImplementedError(
       'In-place operations are no longer supported for Variable.\n'
-      'Use `variable.value //= x`` instead.'
+      'Use `variable[...] //= x`` instead.'
     )
 
   def __imod__(self: V, other) -> V:
     raise NotImplementedError(
       'In-place operations are no longer supported for Variable.\n'
-      'Use `variable.value %= x` instead.'
+      'Use `variable[...] %= x` instead.'
     )
 
   def __ipow__(self: V, other) -> V:
     raise NotImplementedError(
       'In-place operations are no longer supported for Variable.\n'
-      'Use `variable.value **= x`` instead.'
+      'Use `variable[...] **= x`` instead.'
     )
 
   def __ilshift__(self: V, other) -> V:
     raise NotImplementedError(
       'In-place operations are no longer supported for Variable.\n'
-      'Use `variable.value <<= x`` instead.'
+      'Use `variable[...] <<= x`` instead.'
     )
 
   def __irshift__(self: V, other) -> V:
     raise NotImplementedError(
       'In-place operations are no longer supported for Variable.\n'
-      'Use `variable.value >>= x`` instead.'
+      'Use `variable[...] >>= x`` instead.'
     )
 
   def __iand__(self: V, other) -> V:
     raise NotImplementedError(
       'In-place operations are no longer supported for Variable.\n'
-      'Use `variable.value &= x` instead.'
+      'Use `variable[...] &= x` instead.'
     )
 
   def __ixor__(self: V, other) -> V:
     raise NotImplementedError(
       'In-place operations are no longer supported for Variable.\n'
-      'Use `variable.value ^= x` instead.'
+      'Use `variable[...] ^= x` instead.'
     )
 
   def __ior__(self: V, other) -> V:
     raise NotImplementedError(
       'In-place operations are no longer supported for Variable.\n'
-      'Use `variable.value |= x` instead.'
+      'Use `variable[...] |= x` instead.'
     )
 
   __neg__ = _variable_unary_operator('__neg__')

tests/linen/linen_module_test.py

@@ -2315,6 +2315,27 @@ class Foo(nn.Module):
       Foo(1, None)
     Foo(a=1, parent=None)  # type: ignore[call-arg]
 
+  def test_failure_with_sequencelayer(self):
+    # This is a minimal reproducer of the failure seen with
+    # SequenceLayer project and Flax Linen when enabled support for 3.14
+    # See PR: https://github.com/google/flax/pull/5087
+    # Code below is based on
+    # https://github.com/google/flax/pull/5087#issuecomment-3535067361
+    import abc
+    from collections.abc import Iterator
+    from typing import Protocol
+
+    class CheckpointableIterator(Iterator, Protocol):
+      pass
+
+    class Steppable(metaclass=abc.ABCMeta):
+      pass
+
+    isinstance(Steppable, Iterator)
+
+    class SequenceLayer(nn.Module, Steppable):
+      pass
+
   def test_module_path_empty(self):
     rngkey = jax.random.key(0)
     scope = Scope({}, {'params': rngkey}, mutable=['params'])