flax

30b0b0a2 - fix mypy failures due tu numpy update