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,

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'])