Fix signature

PlasmaControl · Feb 29, 2024 · 99ff2e2 · 99ff2e2
1 parent 5663560
commit 99ff2e2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/desc/objectives/objective_funs.py b/desc/objectives/objective_funs.py
@@ -1035,7 +1035,7 @@ def jvp_scaled(self, v, x, constants=None):
         jvpfun = lambda *dx: Derivative.compute_jvp(
             compute_scaled, tuple(range(len(x))), dx, *x
         )
-        sig = "(n)," * (len(x) - 1) + "(n)" + "->(k)"
+        sig = ",".join(f"(n{i})" for i in range(len(x))) + "->(k)"
         return jnp.vectorize(jvpfun, signature=sig)(*v)
 
     def jvp_unscaled(self, v, x, constants=None):
@@ -1061,7 +1061,7 @@ def jvp_unscaled(self, v, x, constants=None):
         jvpfun = lambda *dx: Derivative.compute_jvp(
             compute_unscaled, tuple(range(len(x))), dx, *x
         )
-        sig = "(n)," * (len(x) - 1) + "(n)" + "->(k)"
+        sig = ",".join(f"(n{i})" for i in range(len(x))) + "->(k)"
         return jnp.vectorize(jvpfun, signature=sig)(*v)
 
     def print_value(self, *args, **kwargs):