chore(deps): update dependency jax to ^0.5.0 (#107)

Co-authored-by: renovate[bot] <29139614+renovate[bot]@users.noreply.github.com>

Author: renovate[bot]

poetry.lock

@@ -48,7 +48,7 @@ optional = true
 [tool.poetry.group.benchmark.dependencies]
 numpy = "^2.0.0"
 torch = "^2.5.0"
-jax = {extras = ["cuda12"], version = "^0.4.35"}
+jax = {extras = ["cuda12"], version = "^0.5.0"}
 flax = "^0.10.0"
 
 [tool.poetry.group.docs]