flax
fix transform purity check to avoid tree_map.
#1578
Merged

Loading