Merge branch 'main' of https://www.github.com/google/flax into avoid-…

…units-axes-scan

pytest.ini

@@ -7,11 +7,12 @@ filterwarnings =
     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
 # TODO(#2037): Remove this once Optax releases a new version.
     ignore:jax.tree_util.tree_multimap\(\) is deprecated.*:FutureWarning
-