Fix jax version to 0.6.0 in dependencies [TEMPORARY FIX]

ci_env.yml

@@ -16,7 +16,7 @@ dependencies:
   - pytest-repeat
   - icub-models
   - idyntree >=11.0.0
-  - jax
+  - jax<=0.6.0
   - pytorch
   - jax2torch
   - requests

pyproject.toml

@@ -34,11 +34,11 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-jax = ["jax", "jaxlib"]
+jax = ["jax<=0.6.0", "jaxlib"]
 casadi = ["casadi"]
 pytorch = ["torch", "jax", "jaxlib", "jax2torch"]
 test = [
-  "jax",
+  "jax<=0.6.0",
   "jaxlib",
   "casadi",
   "torch",

src/adam/casadi/inverse_kinematics.py

@@ -582,13 +582,13 @@ def _check_target_type(self, frame: str, expected: TargetType):
     def add_cost(self, cost: cs.MX):
         """Add a custom cost term to the optimization problem.
 
-        This method appends the provided cost term to the `cost_terms` list. 
-        During the `_finalize_problem` step, all cost terms in this list are 
-        combined using `cs.sum(cs.vertcat(*self.cost_terms))` and minimized 
+        This method appends the provided cost term to the `cost_terms` list.
+        During the `_finalize_problem` step, all cost terms in this list are
+        combined using `cs.sum(cs.vertcat(*self.cost_terms))` and minimized
         as part of the objective function.
 
-        Note: This method ensures that the optimization graph is still modifiable 
-        before adding the cost term. Once the problem is finalized (after the 
+        Note: This method ensures that the optimization graph is still modifiable
+        before adding the cost term. Once the problem is finalized (after the
         first call to `solve()`), no new cost terms can be added.
         Args:
             cost (cs.MX): The cost term to add.