flax

0c647cbb - Merge pull request #2025 from jheek:fix-tranform-state-reuse