Python 3.14 compatibility issue

[burnpanck]

System information

• OS Platform and Distribution: macOS Sequoia 15.6.1
• Flax, jax, jaxlib versions: jax==0.7.2 jaxlib==0.7.2 flax==0.12.0
• Python version: 3.14

Note that pip install flax currently typically fails due to tensorstore not yet providing python 3.14 binaries. However, tensorstore effectively is an optional dependency, if one doesn't use orbax-checkpointing either.

Problem you have encountered:

import flax fails, likely due to changes in the handling of annotations (PEP 649 & PEP 749)

What you expected to happen:

Logs, error messages, etc:

Traceback:

>>> import flax
Traceback (most recent call last):
  File "<python-input-0>", line 1, in <module>
    import flax
  File "/Users/yves/.pyenv/versions/3.14.0/Library/Frameworks/Python.framework/Versions/3.14/lib/python3.14/site-packages/flax/__init__.py", line 26, in <module>
    from flax import linen
  File "/Users/yves/.pyenv/versions/3.14.0/Library/Frameworks/Python.framework/Versions/3.14/lib/python3.14/site-packages/flax/linen/__init__.py", line 39, in <module>
    from .activation import (
    ...<26 lines>...
    )
  File "/Users/yves/.pyenv/versions/3.14.0/Library/Frameworks/Python.framework/Versions/3.14/lib/python3.14/site-packages/flax/linen/activation.py", line 21, in <module>
    from flax.linen.module import compact
  File "/Users/yves/.pyenv/versions/3.14.0/Library/Frameworks/Python.framework/Versions/3.14/lib/python3.14/site-packages/flax/linen/module.py", line 3173, in <module>
    class CompactNameScope(Module):
    ...<5 lines>...
        return self.fn(self.module_fn(), *args, **kwargs)
  File "/Users/yves/.pyenv/versions/3.14.0/Library/Frameworks/Python.framework/Versions/3.14/lib/python3.14/site-packages/flax/linen/module.py", line 1042, in __init_subclass__
    cls._customized_dataclass_transform(kw_only)
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^
  File "/Users/yves/.pyenv/versions/3.14.0/Library/Frameworks/Python.framework/Versions/3.14/lib/python3.14/site-packages/flax/linen/module.py", line 1114, in _customized_dataclass_transform
    kw_only_dataclasses.dataclass(
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
      cls,
      ^^^^
    ...<2 lines>...
      extra_fields=extra_fields,
      ^^^^^^^^^^^^^^^^^^^^^^^^^^
    )  # pytype: disable=wrong-keyword-args
    ^
  File "/Users/yves/.pyenv/versions/3.14.0/Library/Frameworks/Python.framework/Versions/3.14/lib/python3.14/site-packages/flax/linen/kw_only_dataclasses.py", line 125, in dataclass
    return wrap if cls is None else wrap(cls)
                                    ~~~~^^^^^
  File "/Users/yves/.pyenv/versions/3.14.0/Library/Frameworks/Python.framework/Versions/3.14/lib/python3.14/site-packages/flax/linen/kw_only_dataclasses.py", line 123, in wrap
    return _process_class(cls, extra_fields=extra_fields, **kwargs)
  File "/Users/yves/.pyenv/versions/3.14.0/Library/Frameworks/Python.framework/Versions/3.14/lib/python3.14/site-packages/flax/linen/kw_only_dataclasses.py", line 191, in _process_class
    cls_annotations = cls.__dict__['__annotations__']
                      ~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^
KeyError: '__annotations__'

Minimal reproducer

• Install python 3.14
• pip install jax msgpack treescope
• pip install --no-deps flax
• python -c "import flax"

[norabelrose]

I'm also running into this problem.

[DiagRisker]

Same here

[vfdev-5]

This branch fixes the annotation issue in flax, however, checkpointing API wont work in python 3.14 until the dependencies like tensorstore release binaries to support it.

[vfdev-5]

Reopening as the #5032 was reverted.

The error in https://github.com/google/sequence-layers to figure out:

File "/build/work/41fc04a72d697b66a8bdbde041e45ac40acd/google3/runfiles/google3/third_party/py/sequence_layers/jax/types.py", line 1297, in <module>
    class SequenceLayer(nn.Module, Steppable):
  File "<frozen abc>", line 106, in __new__
  File "/build/work/41fc04a72d697b66a8bdbde041e45ac40acd/google3/runfiles/google3/third_party/py/flax/linen/module.py", line 1042, in __init_subclass__
    cls._customized_dataclass_transform(kw_only)
  File "/build/work/41fc04a72d697b66a8bdbde041e45ac40acd/google3/runfiles/google3/third_party/py/flax/linen/module.py", line 1114, in _customized_dataclass_transform
    kw_only_dataclasses.dataclass(
  File "/build/work/41fc04a72d697b66a8bdbde041e45ac40acd/google3/runfiles/google3/third_party/py/flax/linen/kw_only_dataclasses.py", line 125, in dataclass
    return wrap if cls is None else wrap(cls)
                                    ^^^^^^^^^
  File "/build/work/41fc04a72d697b66a8bdbde041e45ac40acd/google3/runfiles/google3/third_party/py/flax/linen/kw_only_dataclasses.py", line 123, in wrap
    return _process_class(cls, extra_fields=extra_fields, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/build/work/41fc04a72d697b66a8bdbde041e45ac40acd/google3/runfiles/google3/third_party/py/flax/linen/kw_only_dataclasses.py", line 211, in _process_class
    transformed_cls: type[M] = dataclasses.dataclass(cls, **kwargs)
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/<embedded stdlib>/dataclasses.py", line 1275, in dataclass
  File "/<embedded stdlib>/dataclasses.py", line 1265, in wrap
  File "/<embedded stdlib>/dataclasses.py", line 1018, in _process_class
TypeError: 'parent' is a field but has no type annotation

[cool-RR]

I'm also affected by this problem.

[cool-RR]

Tensorstore just announced that they released Tensorstore 0.1.79 which has wheels for Python 3.14: https://pypi.org/project/tensorstore/#files