flax

b4682076 - Merge pull request #3523 from google:improve-variables-mapping