Merge pull request #513 from meyer-lab/fix-exports-plsr-perf

Miscellaneous housekeeping improvements

.github/workflows/build_doc.yml

@@ -16,7 +16,7 @@ jobs:
       - name: Install Python
         uses: actions/setup-python@v4
         with:
-          python-version: 3.9
+          python-version: '3.10'
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip

.github/workflows/deploy_pypi.yml

@@ -16,7 +16,7 @@ jobs:
       - name: Install Python
         uses: actions/setup-python@v4
         with:
-          python-version: 3.9
+          python-version: '3.10'
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip

.github/workflows/lint.yml

@@ -16,7 +16,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v4
       with:
-        python-version: 3.9
+        python-version: '3.10'
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip

.github/workflows/test.yml

@@ -9,16 +9,20 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        BACKEND: ['numpy', 'pytorch', 'tensorflow', 'jax', 'mxnet']
-        python-version: ['3.10']
+        BACKEND: ['numpy', 'pytorch', 'mxnet']
+        python-version: ['3.11']
         include:
           - BACKEND: 'numpy'
-            python-version: '3.7'
-          - BACKEND: 'numpy'
-            python-version: '3.9'
+            python-version: '3.8'
           - BACKEND: 'numpy'
             python-version: '3.9'
             TENSORLY_TENALG_BACKEND: ['einsum']
+          - BACKEND: 'jax'
+            python-version: '3.11'
+            TENSORLY_TENALG_BACKEND: ['einsum']
+          - BACKEND: 'tensorflow'
+            python-version: '3.11'
+            TENSORLY_TENALG_BACKEND: ['einsum']
 
     steps:
     - uses: actions/checkout@v3

requirements.txt

@@ -2,4 +2,5 @@ numpy
 scipy
 pytest
 pytest-cov
+pytest-randomly
 python_version >= "3.6"

tensorly/__init__.py

@@ -58,16 +58,21 @@
     ndim,
     to_numpy,
     copy,
+    float64,
     concatenate,
     reshape,
     transpose,
     moveaxis,
     arange,
+    any,
     ones,
     zeros,
     zeros_like,
     eye,
     where,
+    conj,
+    index,
+    index_update,
     clip,
     max,
     min,
@@ -88,6 +93,7 @@
     qr,
     kr,
     stack,
+    maximum,
     eps,
     finfo,
     matmul,
@@ -107,6 +113,7 @@
     arccos,
     arctan,
     sinh,
+    sort,
     cosh,
     tanh,
     arcsinh,

tensorly/backend/__init__.py

@@ -1,6 +1,6 @@
 import warnings
 
-from .core import Backend
+from .core import Backend, backend_array
 import importlib
 import os
 import threading
@@ -33,39 +33,28 @@ def __get__(self, instance, cls=None):
 
 class BackendManager(types.ModuleType):
     _functions = [
-        "reshape",
         "moveaxis",
-        "any",
         "trace",
         "shape",
         "ndim",
         "where",
         "copy",
         "transpose",
         "arange",
-        "ones",
-        "zeros",
-        "zeros_like",
         "eye",
         "kron",
         "concatenate",
         "max",
-        "min",
-        "matmul",
         "all",
         "mean",
         "sum",
-        "cumsum",
         "prod",
         "sign",
-        "abs",
-        "sqrt",
         "argmin",
         "argmax",
         "stack",
         "conj",
         "diag",
-        "einsum",
         "log",
         "log2",
         "dot",
@@ -94,7 +83,6 @@ class BackendManager(types.ModuleType):
         "is_tensor",
         "argsort",
         "flip",
-        "count_nonzero",
         "sin",
         "cos",
         "tan",
@@ -115,7 +103,7 @@ class BackendManager(types.ModuleType):
         "atanh",
         "partial_svd",
         "logsumexp",
-    ]
+    ] + backend_array
     _attributes = [
         "int64",
         "int32",

tensorly/backend/core.py

@@ -53,6 +53,7 @@
     "sqrt",
     "abs",
     "min",
+    "maximum",
     "zeros_like",
 ]
 
@@ -501,6 +502,22 @@ def any(tensor, axis=None, keepdims=False, **kwargs):
         """
         return tensor.any(axis=axis, keepdims=keepdims, **kwargs)
 
+    @staticmethod
+    def maximum(x1, x2, *args, **kwargs):
+        """Element-wise maximum of array elements.
+
+        Parameters
+        ----------
+        x1, x2 : tensor
+            The arrays holding the elements to be compared.
+
+        Returns
+        -------
+        tensor
+            The maximum of x1 and x2, element-wise.
+        """
+        raise NotImplementedError
+
     @staticmethod
     def clip(tensor, a_min=None, a_max=None):
         """Clip the values of a tensor to within an interval.
@@ -527,30 +544,36 @@ def clip(tensor, a_min=None, a_max=None):
         raise NotImplementedError
 
     @staticmethod
-    def max(tensor):
+    def max(tensor, axis=None):
         """The max value in a tensor.
 
         Parameters
         ----------
         tensor : tensor
+        axis : int or None, default is None
+            optional, indicates an axis along which to check for non-zero values
 
         Returns
         -------
-        scalar
+        scalar or tensor
+            If axis is None, returns a scalar. Otherwise, returns a tensor of scalars.
         """
         raise NotImplementedError
 
     @staticmethod
-    def min(tensor):
+    def min(tensor, axis=None):
         """The min value in a tensor.
 
         Parameters
         ----------
         tensor : tensor
+        axis : int or None, default is None
+            optional, indicates an axis along which to check for non-zero values
 
         Returns
         -------
-        scalar
+        scalar or tensor
+            If axis is None, returns a scalar. Otherwise, returns a tensor of scalars.
         """
         raise NotImplementedError
 

tensorly/backend/cupy_backend.py

@@ -47,11 +47,6 @@ def ndim(tensor):
     def clip(tensor, a_min=None, a_max=None):
         return cp.clip(tensor, a_min, a_max)
 
-    @staticmethod
-    def lstsq(a, b):
-        x, residuals, _, _ = cp.linalg.lstsq(a, b, rcond=None)
-        return x, residuals
-
     @staticmethod
     def logsumexp(tensor, axis=0):
         max_tensor = cp.max(tensor, axis=axis, keepdims=True)
@@ -94,7 +89,7 @@ def logsumexp(tensor, axis=0):
 ):
     CupyBackend.register_method(name, getattr(cp, name))
 
-for name in ["svd", "qr", "eigh", "solve"]:
+for name in ["svd", "qr", "eigh", "solve", "lstsq"]:
     CupyBackend.register_method(name, getattr(cp.linalg, name))
 
 CupyBackend.register_method("gamma", cp.random.gamma)

tensorly/backend/jax_backend.py

@@ -54,9 +54,8 @@ def ndim(tensor):
         return tensor.ndim
 
     @staticmethod
-    def lstsq(a, b):
-        x, residuals, _, _ = np.linalg.lstsq(a, b, rcond=None, numpy_resid=True)
-        return x, residuals
+    def lstsq(a, b, rcond=None):
+        return np.linalg.lstsq(a, b, rcond=rcond, numpy_resid=True)
 
     @staticmethod
     def logsumexp(tensor, axis=0):

tensorly/backend/mxnet_backend.py

@@ -91,11 +91,6 @@ def svd(self, X, full_matrices=True):
 
         return U, S, V
 
-    @staticmethod
-    def lstsq(a, b):
-        x, residuals, _, _ = np.linalg.lstsq(a, b, rcond=None)
-        return x, residuals
-
     @staticmethod
     def logsumexp(x, axis=0):
         max_x = np.max(x, axis=axis, keepdims=True)
@@ -143,5 +138,5 @@ def logsumexp(x, axis=0):
 ):
     MxnetBackend.register_method(name, getattr(np, name))
 
-for name in ["solve", "qr", "eigh"]:
+for name in ["solve", "qr", "eigh", "lstsq"]:
     MxnetBackend.register_method(name, getattr(np.linalg, name))

tensorly/backend/numpy_backend.py

@@ -33,11 +33,6 @@ def ndim(tensor):
     def clip(tensor, a_min=None, a_max=None):
         return np.clip(tensor, a_min, a_max)
 
-    @staticmethod
-    def lstsq(a, b):
-        x, residuals, _, _ = np.linalg.lstsq(a, b, rcond=None)
-        return x, residuals
-
     @staticmethod
     def logsumexp(tensor, axis=0):
         return scipy.special.logsumexp(tensor, axis=axis)
@@ -77,7 +72,7 @@ def logsumexp(tensor, axis=0):
 ):
     NumpyBackend.register_method(name, getattr(np, name))
 
-for name in ["solve", "qr", "svd", "eigh"]:
+for name in ["solve", "qr", "svd", "eigh", "lstsq"]:
     NumpyBackend.register_method(name, getattr(np.linalg, name))
 
 for name in ["digamma"]:

tensorly/backend/pytorch_backend.py

@@ -190,9 +190,8 @@ def update_index(tensor, index, values):
         tensor.index_put_(index, values)
 
     @staticmethod
-    def lstsq(a, b):
-        x, residuals, _, _ = torch.linalg.lstsq(a, b, rcond=None, driver="gelsd")
-        return x, residuals
+    def lstsq(a, b, rcond=None, driver="gelsd"):
+        return torch.linalg.lstsq(a, b, rcond=rcond, driver=driver)
 
     @staticmethod
     def eigh(tensor):

tensorly/base.py

@@ -1,5 +1,5 @@
+from math import prod
 from . import backend as tl
-from .utils import prod
 
 
 def tensor_to_vec(tensor):

tensorly/decomposition/_cp.py

@@ -408,7 +408,7 @@ def parafac(
             if verbose > 1:
                 print("Mode", mode, "of", tl.ndim(tensor))
 
-            pseudo_inverse = tl.tensor(np.ones((rank, rank)), **tl.context(tensor))
+            pseudo_inverse = tl.ones((rank, rank), **tl.context(tensor))
             for i, factor in enumerate(factors):
                 if i != mode:
                     pseudo_inverse = pseudo_inverse * tl.dot(
@@ -426,8 +426,6 @@ def parafac(
                 tl.solve(tl.conj(tl.transpose(pseudo_inverse)), tl.transpose(mttkrp))
             )
             factors[mode] = factor
-            if normalize_factors and mode != modes_list[-1]:
-                weights, factors = cp_normalize((weights, factors))
 
         # Will we be performing a line search iteration
         if linesearch and iteration % 2 == 0 and iteration > 5:

tensorly/decomposition/_tr.py

@@ -216,7 +216,7 @@ def tensor_ring_als(
 
             if ls_solve == "lstsq":
                 # Solve least squares problem directly
-                sol, _ = tl.lstsq(design_mat, tensor_unf)
+                sol, *_ = tl.lstsq(design_mat, tensor_unf)
 
             elif ls_solve == "normal_eq":
                 # Solve least squares problem via normal equations

tensorly/decomposition/tests/test_cmtf_als.py

@@ -10,8 +10,8 @@
 def test_coupled_matrix_tensor_3d_factorization():
     I = 21
     J = 12
-    K = 8
-    M = 7
+    K = 16
+    M = 14
     R = 3
 
     tensor_cp_true = random_cp((I, J, K), rank=R, normalise_factors=False)

tensorly/metrics/entropy.py

@@ -2,7 +2,6 @@
 import tensorly as tl
 from .. import backend as T
 from ..tt_tensor import tt_to_tensor
-from ..utils import prod
 
 # Authors: Taylor Lee Patti <taylorpatti@g.harvard.edu>
 #          Jean Kossaifi
@@ -25,7 +24,7 @@ def vonneumann_entropy(tensor):
     where p_i are the probabilities that each state is occupied
     (the eigenvalues of the density matrix).
     """
-    square_dim = int(math.sqrt(prod(tensor.shape)))
+    square_dim = int(math.sqrt(math.prod(tensor.shape)))
     tensor = tl.reshape(tensor, (square_dim, square_dim))
     try:
         eig_vals = T.eigh(tensor)[0]