flax

upgrade python to 3.10 + use pyupgrade

#4038

Merged