enable tensorsolve

pytorch · Oct 18, 2024 · 9457cc0 · 9457cc0
1 parent 2d73a5f
commit 9457cc0
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 3 deletions.
diff --git a/experimental/torch_xla2/test/test_ops.py b/experimental/torch_xla2/test/test_ops.py
@@ -32,7 +32,6 @@
     "linalg.lu_solve",
     "linalg.matrix_norm",
     "linalg.matrix_power",
-    "linalg.tensorsolve",
     "masked.median",
     "max_pool2d_with_indices_backward",
     "nn.functional.adaptive_avg_pool3d",
@@ -240,7 +239,7 @@ def test_reference_eager(self, device, dtype, op):
         continue
       check_output = op.name not in random_ops
 
-      #print("[DEBUG] sample_input: ", sample_input)
+      print("[DEBUG] sample_input: ", sample_input)
 
       # TODO: this is a workaround to skip int64 cast for linspace
       # reference: https://github.com/pytorch/xla/issues/7505#issuecomment-2400895692 and subsequent comments

diff --git a/experimental/torch_xla2/torch_xla2/ops/jtorch.py b/experimental/torch_xla2/torch_xla2/ops/jtorch.py
@@ -321,3 +321,8 @@ def lu_solve(b, LU_data, LU_pivots, **kwargs):
   _pivots = LU_pivots - 1
   x = jax.scipy.linalg.lu_solve((LU_data, _pivots), b)
   return x
+
+@register_function(torch.linalg.tensorsolve)
+def linalg_tensorsolve(a, b):
+  res, _ = jaten._aten__linalg_solve_ex(a, b)
+  return res
diff --git a/experimental/torch_xla2/torch_xla2/tensor.py b/experimental/torch_xla2/torch_xla2/tensor.py
@@ -209,7 +209,8 @@ def __torch_function__(self,
         if len(args) >= 2 and type(args[1]) == int:
           if ((args[1])%4 == 0):
             return args[0]
-      return func(*args, **(kwargs or {}))
+      ret = func(*args, **(kwargs or {}))
+      return ret
 
 
 class XLADispatchMode(torch_dispatch.TorchDispatchMode):