flax

Allow custom axis metadata annotation during transforms

#4637

Merged