From 72fda76ed910d81f9704c19f625818c3e52aa253 Mon Sep 17 00:00:00 2001
From: Anish Karthik <89824626+anishfish2@users.noreply.github.com>
Date: Tue, 24 Sep 2024 18:31:55 -0700
Subject: [PATCH] Avoids error removing linalg.eig from skiplist by modifying
 atol and rtol (#8068)

---
 experimental/torch_xla2/test/test_ops.py        | 3 +--
 experimental/torch_xla2/torch_xla2/ops/jaten.py | 2 +-
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/experimental/torch_xla2/test/test_ops.py b/experimental/torch_xla2/test/test_ops.py
index 3185d996b26..1c67da372c1 100644
--- a/experimental/torch_xla2/test/test_ops.py
+++ b/experimental/torch_xla2/test/test_ops.py
@@ -42,7 +42,6 @@
     "linalg.cholesky_ex",
     "linalg.cond",
     "linalg.det",
-    "linalg.eig",
     "linalg.eigh",
     "linalg.eigvalsh",
     "linalg.householder_product",
@@ -184,7 +183,7 @@
   'nn.functional.feature_alpha_dropout',
 }
 
-atol_dict = {"matrix_exp": (2e-1, 2e-4), "linalg.pinv": (8e-1, 2e0)}
+atol_dict = {"matrix_exp": (2e-1, 2e-4), "linalg.pinv": (8e-1, 2e0), "linalg.eig":(2e0, 3e0)}
 
 def diff_output(testcase, output1, output2, rtol, atol, equal_nan=True, check_output=True):
   if isinstance(output1, torch.Tensor):
diff --git a/experimental/torch_xla2/torch_xla2/ops/jaten.py b/experimental/torch_xla2/torch_xla2/ops/jaten.py
index 90adbb9cc13..ac56cc03fa9 100644
--- a/experimental/torch_xla2/torch_xla2/ops/jaten.py
+++ b/experimental/torch_xla2/torch_xla2/ops/jaten.py
@@ -2196,7 +2196,7 @@ def _aten_igamma(input, other):
 
 @op(torch.ops.aten.linalg_eig)
 def _aten_linalg_eig(A):
-  return jax.numpy.linalg.eig(A)
+  return jnp.linalg.eig(A)
 
 
 # aten.lcm