Merge pull request tensorly#589 from JeanKossaifi/main

FIX partial_unfold + add test

Author: JeanKossaifi

tensorly/base.py

@@ -113,7 +113,7 @@ def partial_unfold(tensor, mode=0, skip_begin=1, skip_end=0, ravel_tensors=False
         new_shape = [tensor.shape[i] for i in range(skip_begin)] + new_shape
 
     if skip_end:
-        new_shape += [tensor.shape[-i] for i in range(1, 1 + skip_end)]
+        new_shape += [tensor.shape[-i] for i in range(skip_end, 0, -1)]
 
     return tl.reshape(tl.moveaxis(tensor, mode + skip_begin, skip_begin), new_shape)
 

tensorly/tests/test_core.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from .. import backend as T
+import tensorly as T
 from ..base import fold, unfold
 from ..base import partial_fold, partial_unfold
 from ..base import tensor_to_vec, vec_to_tensor
@@ -320,6 +320,12 @@ def test_partial_tensor_to_vec():
     for j in range(n_samples):  # test for each sample
         assert_array_equal(T.transpose(vectorised)[j], vec_X + j)
 
+    tensor = T.randn((2, 3, 4, 5))
+    TT = partial_tensor_to_vec(tensor, skip_begin=0, skip_end=2)
+    assert T.shape(TT) == (6, 4, 5)
+    rec = partial_vec_to_tensor(TT, T.shape(tensor), skip_begin=0, skip_end=2)
+    assert T.shape(rec) == T.shape(tensor)
+
 
 def test_partial_vec_to_tensor():
     """Test for partial_vec_to_tensor"""