Merge pull request #2715 from marcvanzee:tree_multimap

PiperOrigin-RevId: 495268625

Author: Flax Authors

pytest.ini

@@ -5,20 +5,12 @@ filterwarnings =
     error
 # Jax warning when no gpu/tpu found.
     ignore:No GPU/TPU found, falling back to CPU.*:UserWarning
-# Jax warns about XLA not being able to use donated buffers.
-    ignore:Some donated buffers were not usable.*:UserWarning
 # We should remove our dependence on jtu.
     ignore:jax.test_util.check_eq is deprecated and will soon be removed.*:FutureWarning
 # Tensorflow's fast_tensor_util.pyx cython raises:
 # ImportWarning: can't resolve package from __spec__ or __package__, falling back on __name__ and __path__
     ignore:can't resolve package from.*:ImportWarning
-# Jax imports flatbuffers which imports imp in a compat file.
-    ignore:the imp module is deprecated.*:DeprecationWarning
 # traverse_util.Traversal will be removed soon.
     ignore:`flax.traverse_util.Traversal` will be deprecated.*:DeprecationWarning
-# TODO: Will revisit all the deprecation warnings next week.
-    ignore:jax.tree_.*:FutureWarning
-# TF internal deprecation warnings.
-    ignore:Call to deprecated create function .*Descriptor.*:DeprecationWarning
 # distutils DeprecationWarning
     ignore:distutils Version classes are deprecated.*:DeprecationWarning