Merge pull request #5032 from google:fix-py314-kw-only-dataclasses

PiperOrigin-RevId: 822149189

Author: Flax Authors

.github/workflows/flax_test.yml

@@ -63,7 +63,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ['3.11', '3.12', '3.13']
+        python-version: ['3.11', '3.12', '3.13', '3.14']
     steps:
     - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
     - name: Set up Python ${{ matrix.python-version }}
@@ -159,3 +159,30 @@ jobs:
            "description": "'$status'",
            "context": "github-actions/Build"
            }'
+
+  # This is a temporary workflow to test flax on Python 3.14 and
+  # skipping deps like tensorstore, tensorflow etc
+  tests-python314:
+    name: Run Tests on Python 3.14
+    needs: [pre-commit, commit-count]
+    runs-on: ubuntu-24.04-16core
+    steps:
+    - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+    - name: Setup uv
+      uses: astral-sh/setup-uv@887a942a15af3a7626099df99e897a18d9e5ab3a # v5.1.0
+      with:
+        version: "0.9.2"
+        python-version: "3.14"
+        activate-environment: true
+        enable-cache: true
+
+    - name: Install dependencies
+      run: |
+        uv pip install jax msgpack treescope rich typing_extensions PyYAML optax cloudpickle
+        uv pip install pytest pytest-custom_exit_code pytest-xdist pytest-cov
+        uv pip install -e . --no-deps
+    - name: Test with pytest
+      run: |
+        export XLA_FLAGS='--xla_force_host_platform_device_count=4'
+        find tests/ -name "*.py" | grep -vE 'checkpoint|integr|io|tensorboard' | xargs pytest -n auto
+

flax/linen/kw_only_dataclasses.py

@@ -127,7 +127,7 @@ def wrap(cls):
 
 def _process_class(cls: type[M], extra_fields=None, **kwargs):
   """Transforms `cls` into a dataclass that supports kw_only fields."""
-  if '__annotations__' not in cls.__dict__:
+  if not hasattr(cls, "__annotations__"):
     cls.__annotations__ = {}
 
   # The original __dataclass_fields__ dicts for all base classes.  We will
@@ -174,7 +174,7 @@ def _process_class(cls: type[M], extra_fields=None, **kwargs):
   for base in reversed(cls.__mro__[1:]):
     if not dataclasses.is_dataclass(base):
       continue
-    base_annotations = base.__dict__.get('__annotations__', {})
+    base_annotations = inspect.get_annotations(base)
     base_dataclass_fields[base] = dict(
       getattr(base, '__dataclass_fields__', {})
     )
@@ -188,7 +188,8 @@ def _process_class(cls: type[M], extra_fields=None, **kwargs):
         del base.__dataclass_fields__[field_name]
 
   # Remove any keyword-only fields from this class.
-  cls_annotations = cls.__dict__['__annotations__']
+  # Note: in Python 3.14+ cls.__dict__ does not contain __annotation__ key but __annotations_cache__
+  cls_annotations = getattr(cls, "__annotations__")
   for name, annotation in list(cls_annotations.items()):
     value = getattr(cls, name, None)
     if (

flax/linen/module.py

@@ -1061,7 +1061,7 @@ def _customized_dataclass_transform(cls, kw_only: bool):
     3. Generate a hash function (if not provided by cls).
     """
     # Check reserved attributes have expected type annotations.
-    annotations = dict(cls.__dict__.get('__annotations__', {}))
+    annotations = inspect.get_annotations(cls)
     if annotations.get('parent', _ParentType) != _ParentType:
       raise errors.ReservedModuleAttributeError(annotations)
     if annotations.get('name', str) not in ('str', str, Optional[str]):

pyproject.toml

@@ -17,8 +17,8 @@ dependencies = [
     "jax>=0.7.1",
     "msgpack",
     "optax",
-    "orbax-checkpoint",
-    "tensorstore",
+    "orbax-checkpoint; python_version<'3.14'",  # temporary skip orbax-checkpoint for py3.14
+    "tensorstore; python_version<'3.14'",  # temporary skip tensorstore for py3.14
     "rich>=11.1",
     "typing_extensions>=4.2",
     "PyYAML>=5.4.1",