Python 3.14 compatibility issue #5027
Description
System information
- OS Platform and Distribution: macOS Sequoia 15.6.1
- Flax, jax, jaxlib versions:
jax==0.7.2jaxlib==0.7.2flax==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 treescopepip install --no-deps flaxpython -c "import flax"