Repository: tensorly/tensorly

Issue:
make_svd_non_negative only returns the updated U matrix
#### Describe the bug
I'm not sure if this is a bug or if this is actually how this function is supposed to behave but currently, it only returns the updated U value. Incase it's not a bug, I don't understand why the H matrix is computed in this function (https://github.com/tensorly/tensorly/blob/6c7ecfbfbfc088cb1a423e6a5970e19565844651/tensorly/tenalg/svd.py#L65)

Looking forward to your response.

Discussion:

Patch:
diff --git a/tensorly/tenalg/svd.py b/tensorly/tenalg/svd.py
index e9f1767..212ccb2 100644
--- a/tensorly/tenalg/svd.py
+++ b/tensorly/tenalg/svd.py
@@ -87,7 +87,7 @@ def make_svd_non_negative(tensor, U, S, V, nntype=True):
     W = tl.index_update(W, tl.index[:, 0], tl.sqrt(S[0]) * tl.abs(U[:, 0]))
     H = tl.index_update(H, tl.index[0, :], tl.sqrt(S[0]) * tl.abs(V[0, :]))
 
-    for j in range(1, tl.shape(U)[1]):
+    for j in range(1, min(tl.shape(U)[1], tl.shape(V)[0])):
         x, y = U[:, j], V[j, :]
 
         # extract positive and negative parts of column vectors
@@ -119,15 +119,17 @@ def make_svd_non_negative(tensor, U, S, V, nntype=True):
 
     if nntype == "nndsvd":
         W = soft_thresholding(W, eps)
+        H = soft_thresholding(H, eps)
     elif nntype == "nndsvda":
         avg = tl.mean(tensor)
         W = tl.where(W < eps, tl.ones(tl.shape(W), **tl.context(W)) * avg, W)
+        H = tl.where(H < eps, tl.ones(tl.shape(H), **tl.context(H)) * avg, H)
     else:
         raise ValueError(
             f'Invalid nntype parameter: got {nntype} instead of one of ("nndsvd", "nndsvda")'
         )
 
-    return W
+    return W, H
 
 
 def randomized_range_finder(A, n_dims, n_iter=2, random_state=None):
@@ -425,13 +427,18 @@ def svd_interface(
 
     if mask is not None and n_eigenvecs is not None:
         for _ in range(n_iter_mask_imputation):
-            matrix = matrix * mask + (U @ tl.diag(S) @ V) * (1 - mask)
+            # Workaround to avoid needing fill_diagonal
+            St = tl.eye(tl.shape(U)[1], tl.shape(V)[0], **tl.context(matrix))
+            for i in range(tl.shape(S)[0]):
+                St = tl.index_update(St, tl.index[i, i], S[i])
+
+            matrix = matrix * mask + (U @ St @ V) * (1 - mask)
             U, S, V = svd_fun(matrix, n_eigenvecs=n_eigenvecs, **kwargs)
 
     if flip_sign:
         U, V = svd_flip(U, V, u_based_decision=u_based_flip_sign)
 
     if non_negative is not False and non_negative is not None:
-        U = make_svd_non_negative(matrix, U, S, V, non_negative)
+        U, V = make_svd_non_negative(matrix, U, S, V, non_negative)
 
     return U, S, V


Test patch:
diff --git a/tensorly/tenalg/tests/test_svd.py b/tensorly/tenalg/tests/test_svd.py
new file mode 100644
index 0000000..04dd91e
--- /dev/null
+++ b/tensorly/tenalg/tests/test_svd.py
@@ -0,0 +1,30 @@
+import pytest
+from ...testing import assert_
+from ..svd import svd_interface
+import tensorly as tl
+from ...testing import assert_
+
+
+@pytest.mark.parametrize("shape", [(10, 5), (10, 10), (5, 10)])
+@pytest.mark.parametrize("rank", [1, 3, 5, 8, 10])
+@pytest.mark.parametrize("nn", [False, True])
+def test_svd_interface_shape(shape, rank, nn):
+    """Test that the SVD interface handles shape edge cases."""
+    rng = tl.check_random_state(1234)
+
+    X = tl.tensor(rng.random_sample(shape))
+    X_mask = tl.zeros_like(X)
+    X_mask = tl.index_update(X_mask, tl.index[1, 3], 1.0)
+
+    U, S, V = svd_interface(
+        X,
+        n_eigenvecs=rank,
+        non_negative=nn,
+        mask=X_mask,
+        n_iter_mask_imputation=3,
+    )
+    assert_(tl.all(S >= 0.0))
+
+    if nn:
+        assert_(tl.all(U >= 0.0))
+        assert_(tl.all(V >= 0.0))