flax

2970c6ff - Merge pull request #5237 from mohsinm-dev:fix/pytreenode-generic-parameters