Fix op info test for svd (pytorch#8289)

dvhg · Oct 21, 2024 · b8ee2c2 · b8ee2c2
1 parent b2e0e33
commit b8ee2c2
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/experimental/torch_xla2/test/test_ops.py b/experimental/torch_xla2/test/test_ops.py
@@ -84,7 +84,6 @@
     "special.scaled_modified_bessel_k1",
     "special.spherical_bessel_j0",
     "special.zeta",
-    "svd",
     "svd_lowrank",
     "unfold_copy",
     "unfold",
@@ -136,6 +135,7 @@
              "linalg.eigvalsh": (5e1, 3e0),
              "linalg.pinv": (8e-1, 2e0),
              "linalg.svd": (1e0, 1e0),
+             "svd": (1e0, 1e0),
              "matrix_exp": (2e-1, 2e-4),
              "cdist": (5e1, 3e0)}
 

diff --git a/experimental/torch_xla2/torch_xla2/ops/jtorch.py b/experimental/torch_xla2/torch_xla2/ops/jtorch.py
@@ -355,6 +355,16 @@ def linalg_solve_ex(a, b):
 def linalg_svd(a, full_matrices=True, **kwargs):
   return jaten._aten__linalg_svd(a, full_matrices=full_matrices, **kwargs)
 
+@register_function(torch.svd)
+def svd(a, some=True, compute_uv=True):
+  if not compute_uv:
+    S = jaten._aten__linalg_svd(a, full_matrices=False)[1]
+    U = jnp.zeros((a.shape[-2], a.shape[-2]), dtype=a.dtype)
+    V = jnp.zeros((a.shape[-1], a.shape[-1]), dtype=a.dtype)
+    return U, S, V
+  U, S, V = jaten._aten__linalg_svd(a, full_matrices=not some)
+  return U, S, jnp.matrix_transpose(V)
+
 @register_function(torch.cdist)
 def _cdist(x1, x2, p=2.0, compute_mode='use_mm_for_euclid_dist_if_necessary'):
     return jaten._aten_cdist(x1, x2, p, compute_mode)