Copybara import of the project:

--
6622efc by vfdev-5 <[email protected]>:

Support for python 3.14

Adapted the way to access `__annotations__` object.
Flax linen code relies on `cls.__dict__["__annotations__"]`, but
in Python 3.14+ `cls.__dict__` does not contain `__annotation__` key anymore.

Fixes #5027

PiperOrigin-RevId: 822202664

Author: Flax Team

.github/workflows/flax_test.yml

@@ -63,7 +63,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ['3.11', '3.12', '3.13', '3.14']
+        python-version: ['3.11', '3.12', '3.13']
     steps:
     - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
     - name: Set up Python ${{ matrix.python-version }}
@@ -159,30 +159,3 @@ 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 not hasattr(cls, "__annotations__"):
+  if '__annotations__' not in cls.__dict__:
     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 = inspect.get_annotations(base)
+    base_annotations = base.__dict__.get('__annotations__', {})
     base_dataclass_fields[base] = dict(
       getattr(base, '__dataclass_fields__', {})
     )
@@ -188,8 +188,7 @@ def _process_class(cls: type[M], extra_fields=None, **kwargs):
         del base.__dataclass_fields__[field_name]
 
   # Remove any keyword-only fields from this class.
-  # Note: in Python 3.14+ cls.__dict__ does not contain __annotation__ key but __annotations_cache__
-  cls_annotations = getattr(cls, "__annotations__")
+  cls_annotations = cls.__dict__['__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 = inspect.get_annotations(cls)
+    annotations = dict(cls.__dict__.get('__annotations__', {}))
     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; python_version<'3.14'",  # temporary skip orbax-checkpoint for py3.14
-    "tensorstore; python_version<'3.14'",  # temporary skip tensorstore for py3.14
+    "orbax-checkpoint",
+    "tensorstore",
     "rich>=11.1",
     "typing_extensions>=4.2",
     "PyYAML>=5.4.1",