diff --git a/pyproject.toml b/pyproject.toml
index e486430064c298cb3162945bd694a790e711753e..3ce5cc8b5391c0843c2df7bae7ff52d6c913193d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -48,7 +48,6 @@ dependencies = [
 [project.optional-dependencies]
 all = [  # all non-tests extra dependencies
     "dm-haiku == 0.0.9",
-    "jax == 0.4.4",
     "functorch",
     "grpcio >= 1.45",
     "jax[cpu] == 0.4.4",
@@ -67,8 +66,8 @@ grpc = [
     "protobuf >= 3.19",
 ]
 haiku = [
-    "jax == 0.4.4",
     "dm-haiku == 0.0.9",
+    "jax[cpu] == 0.4.4",
 ]
 tensorflow = [
     "tensorflow ~= 2.5",