flax

fix kw_only_dataclasses for python 3.14 (part 2)

#5135

Merged