diff --git a/experimental/torch_xla2/test/test_ops.py b/experimental/torch_xla2/test/test_ops.py index aabd2a11352..f64c319be3c 100644 --- 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 index e355094694d..b553d5e4b4d 100644 --- 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 \ No newline at end of file diff --git a/experimental/torch_xla2/torch_xla2/tensor.py b/experimental/torch_xla2/torch_xla2/tensor.py index 6e446fb3874..7fdb8da6024 100644 --- 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):