flax

3609a43f - Merge pull request #451 from jheek:unified-variable-axes