Update minimal JAX version to latest (0.3.16).

0.3.16 includes jax-ml/jax#11807, which supports more flexible sharding constraints under vmap (see also flax/linen/partitioning.py).

0.3.16 removes jax.tree_util.tree_multimap, so we are also updating to the latest jgraph version (in which tree_multimap has been replaced).

PiperOrigin-RevId: 467265619

pytest.ini

@@ -14,8 +14,6 @@ filterwarnings =
     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 Jgraph uses jax.tree_map.
-    ignore:jax.tree_util.tree_multimap\(\) is deprecated.*:FutureWarning
 # 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. 

setup.py

@@ -26,7 +26,7 @@
 
 install_requires = [
     "numpy>=1.12",
-    "jax>=0.3.15",
+    "jax>=0.3.16",
     "matplotlib",  # only needed for tensorboard export
     "msgpack",
     "optax",
@@ -40,7 +40,7 @@
     "clu",  # All examples.
     "gym==0.18.3",
     "jaxlib",
-    "jraph",
+    "jraph>=0.0.6dev0",
     "ml-collections",
     "opencv-python",
     "pytest",