Add support for erfinv (#8081)

pytorch · Sep 27, 2024 · ed68381 · ed68381
1 parent f088810
commit ed68381
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/experimental/torch_xla2/test/test_ops.py b/experimental/torch_xla2/test/test_ops.py
@@ -28,7 +28,6 @@
     "diagonal_copy",
     "diagonal_scatter",
     "digamma",
-    "erfinv",
     "exponential",
     "gcd",
     "geometric",

diff --git a/experimental/torch_xla2/torch_xla2/ops/jaten.py b/experimental/torch_xla2/torch_xla2/ops/jaten.py
@@ -2057,6 +2057,12 @@ def _aten_erf(x):
   return jax.lax.erf(x)
 
 
+@op(torch.ops.aten.erfinv)
+@op_base.promote_int_input
+def _aten_erfinv(input):
+  return jax.lax.erf_inv(input)
+
+
 # aten.exp
 @op(torch.ops.aten.exp)
 def _aten_exp(input):