Op info test for logdet, linalg.slogdet (#7503) (#7958) (#8004)

experimental/torch_xla2/test/test_ops.py

@@ -98,7 +98,6 @@
     "linalg.multi_dot",
     "linalg.norm",
     "linalg.pinv",
-    "linalg.slogdet",
     "linalg.solve",
     "linalg.solve_ex",
     "linalg.solve_triangular",
@@ -110,7 +109,6 @@
     "linalg.vector_norm",
     "linspace",
     "log_normal",
-    "logdet",
     "logspace",
     "lu",
     "lu_solve",

experimental/torch_xla2/torch_xla2/ops/jaten.py

@@ -3766,6 +3766,13 @@ def _aten_linalg_matrix_exp(input):
   return jax.scipy.linalg.expm(input)
 
 
+# torch._linalg.slogdet
+@op(torch.ops.aten._linalg_slogdet)
+def _aten__linalg_slogdet(input):
+  res = jnp.linalg.slogdet(input)
+  return res.sign, res.logabsdet
+
+
 @op(torch.ops.aten.median)
 def _aten_median(self, dim=None, keepdim=False):
   output = _with_reduction_scalar(functools.partial(jnp.quantile, q=0.5, method='lower'), self, dim=dim, keepdim=keepdim).astype(self.dtype)

experimental/torch_xla2/torch_xla2/ops/jtorch.py

@@ -191,6 +191,7 @@ def _sparse_mm(mat1, mat2, reduce='sum'):
 def _aten_isclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False):
   return jnp.isclose(input, other, rtol, atol, equal_nan)
 
+
 @register_function(torch.ones)
 def _ones(*size: int, dtype=None, **kwargs):
   if len(size) == 1 and isinstance(size[0], collections.abc.Iterable):
@@ -265,3 +266,15 @@ def randint(
   *args, **kwargs
 ):
   return torch.ops.aten.randint(*args, **kwargs)
+
+
+@register_function(torch.logdet)
+def logdet(input):
+  _, logabsdet = jaten._aten__linalg_slogdet(input)
+  return logabsdet
+
+
+@register_function(torch.linalg.slogdet)
+def slogdet(input):
+  sign, logabsdet = jaten._aten__linalg_slogdet(input)
+  return torch.return_types.slogdet((sign, logabsdet))