upgrade jax

pyproject.toml

@@ -8,9 +8,8 @@ readme = "README.md"
 
 [tool.poetry.dependencies]
 python = "3.10.12"
-ruff = "0.4.2"
-jax = "0.4.26"
-jaxlib = "0.4.26"
+jax = "0.4.28"
+jaxlib = "0.4.28"
 ray = {version = "2.20.0", extras = ["default", "data"]}
 flax = "0.8.3"
 optax = "0.2.2"
@@ -24,18 +23,35 @@ rich = "13.7.1"
 pydantic = "2.7.1"
 jmp = "0.0.4"
 jaxtyping = "0.2.28"
-pre-commit = "3.7.0"
-ipdb = "0.13.13"
 beartype = "0.18.5"
 tensorflow = "2.16.1"
 tensorflow-datasets = "4.9.4"
 pillow = "10.3.0"
-coverage = "7.5.1"
 wandb = "0.17.0"
 ml-collections = "0.1.1"
 
 
 
+[[tool.poetry.source]]
+name = "jax_tpu"
+url = "https://storage.googleapis.com/jax-releases/libtpu_releases.html"
+priority = "supplemental"
+
+[tool.poetry.group.tpu]
+optional = true
+
+[tool.poetry.group.tpu.dependencies]
+libtpu-nightly = { version="0.1.dev20240521", source="jax_tpu" }
+
+[tool.poetry.group.dev]
+optional = true
+
+[tool.poetry.group.dev.dependencies]
+ruff = "0.4.2"
+pre-commit = "3.7.0"
+ipdb = "0.13.13"
+coverage = "7.5.1"
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"