Infinite loop within tensor cross

[ryanlevy]

Describe the bug

For certain rank inputs to tensor_train_cross the algorithm gets stuck in a while loop.

Steps or Code to Reproduce

import numpy as np
import tensorly as tl
from tensorly.contrib.decomposition import tensor_train_cross

tensor = tl.tensor(np.arange(2**6).reshape(2,2,2,2,2,2))
factors = tensor_train_cross(tensor, [1,3,3,3,3,3,1])

---------------------------------------------------------------------------
KeyboardInterrupt                         Traceback (most recent call last)
...
File /home/rlevy/venvs/tensors/lib/python3.9/site-packages/tensorly/contrib/decomposition/_tt_cross.py:127, in <listcomp>(.0)
    125         newidx = tuple([rng.randint(tensor_shape[j]) for j in range(k_col_idx + 1, tensor_order)])
    126         while newidx in col_idx[k_col_idx]:
--> 127             newidx = tuple([rng.randint(tensor_shape[j]) for j in range(k_col_idx + 1, tensor_order)])
    129         col_idx[k_col_idx].append(newidx)
    131 # Initialize the cores of tensor-train

Versions

Linux-4.18.0-372.19.1.el8_6.x86_64-x86_64-with-glibc2.28
Python 3.9.12 (main, Jun 16 2022, 18:08:55) 
[GCC 10.3.0]
NumPy 1.21.6
SciPy 1.8.0
TensorLy 0.7.0

[meandmytram]

Second the bug!

[JeanKossaifi]

Thanks for reporting @meandmytram @ryanlevy! Would you be able to open a small PR to fix the issue?

[meandmytram]

Hey @JeanKossaifi, thanks for a quick response! While normally I'd be willing to help, now I'm super busy so I'm afraid I simply won't have enough time to dig in the source code of the algorithm.