fix the custom jvp for safe_arcsin (#484)

google · Jun 7, 2024 · 2b95fe9 · 2b95fe9
1 parent 69539e9
commit 2b95fe9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/brax/math.py b/brax/math.py
@@ -216,7 +216,7 @@ def safe_arcsin(x: jax.Array) -> jax.Array:
 def _safe_arcsin_jvp(primal, tangent):
   (x,) = primal
   (x_dot,) = tangent
-  primal_out = safe_arccos(x)
+  primal_out = safe_arcsin(x)
   tangent_out = x_dot / jp.sqrt(1.0 - jp.clip(x, -1 + 1e-7, 1 - 1e-7) ** 2.0)
   return primal_out, tangent_out