flax

added `tree_map` deprecation warning filter

#3828

Merged