flax

f2094bac - Merge pull request #2831 from cgarciae:improve-map_variables-v2