Revert backend changes

meyer-lab · Aug 12, 2024 · b6c6339 · b6c6339
1 parent d64a88b
commit b6c6339
Show file tree

Hide file tree

Showing 2 changed files with 85 additions and 62 deletions.
diff --git a/tensorly/backend/__init__.py b/tensorly/backend/__init__.py
@@ -1,6 +1,6 @@
 import warnings
 
-from .core import Backend, backend_array, backend_types, backend_basic_math
+from .core import Backend, backend_array
 import importlib
 import os
 import threading
@@ -32,71 +32,93 @@ def __get__(self, instance, cls=None):
 
 
 class BackendManager(types.ModuleType):
-    _functions = (
-        [
-            "moveaxis",
-            "trace",
-            "shape",
-            "ndim",
-            "where",
-            "copy",
-            "transpose",
-            "arange",
-            "eye",
-            "kron",
-            "concatenate",
-            "max",
-            "all",
-            "mean",
-            "sum",
-            "prod",
-            "sign",
-            "argmin",
-            "argmax",
-            "stack",
-            "conj",
-            "diag",
-            "dot",
-            "tensordot",
-            "clip",
-            "kron",
-            "diagonal",
-            "lstsq",
-            "eps",
-            "finfo",
-            "solve",
-            "svd",
-            "qr",
-            "randn",
-            "gamma",
-            "digamma",
-            "check_random_state",
-            "sort",
-            "eigh",
-            "index_update",
-            "context",
-            "tensor",
-            "norm",
-            "to_numpy",
-            "is_tensor",
-            "argsort",
-            "flip",
-            "asin",
-            "acos",
-            "atan",
-            "asinh",
-            "acosh",
-            "atanh",
-            "logsumexp",
-        ]
-        + backend_array
-        + backend_basic_math
-    )
+    _functions = [
+        "moveaxis",
+        "trace",
+        "shape",
+        "ndim",
+        "where",
+        "copy",
+        "transpose",
+        "arange",
+        "eye",
+        "kron",
+        "concatenate",
+        "max",
+        "all",
+        "mean",
+        "sum",
+        "prod",
+        "sign",
+        "argmin",
+        "argmax",
+        "stack",
+        "conj",
+        "diag",
+        "log",
+        "log2",
+        "dot",
+        "tensordot",
+        "diagonal",
+        "exp",
+        "clip",
+        "kr",
+        "kron",
+        "lstsq",
+        "eps",
+        "finfo",
+        "solve",
+        "svd",
+        "qr",
+        "randn",
+        "gamma",
+        "digamma",
+        "check_random_state",
+        "sort",
+        "eigh",
+        "index_update",
+        "context",
+        "tensor",
+        "norm",
+        "to_numpy",
+        "is_tensor",
+        "argsort",
+        "flip",
+        "sin",
+        "cos",
+        "tan",
+        "asin",
+        "acos",
+        "atan",
+        "arcsin",
+        "arccos",
+        "arctan",
+        "sinh",
+        "cosh",
+        "tanh",
+        "arcsinh",
+        "arccosh",
+        "arctanh",
+        "asinh",
+        "acosh",
+        "atanh",
+        "partial_svd",
+        "logsumexp",
+    ] + backend_array
     _attributes = [
+        "int64",
+        "int32",
+        "float64",
+        "float32",
+        "pi",
+        "e",
+        "inf",
         "nan",
+        "complex128",
+        "complex64",
         "index",
         "backend_name",
-    ] + backend_types
+    ]
     available_backend_names = ["numpy", "pytorch", "tensorflow", "cupy", "jax"]
     _default_backend = "numpy"
     _loaded_backends = dict()

diff --git a/tensorly/contrib/sparse/backend/__init__.py b/tensorly/contrib/sparse/backend/__init__.py
@@ -108,6 +108,7 @@ def inner(*args, **kwargs):
 norm = dispatch_sparse(backend.norm)
 dot = dispatch_sparse(backend.dot)
 kron = dispatch_sparse(backend.kron)
+kr = dispatch_sparse(backend.kr)
 solve = dispatch_sparse(backend.solve)
 qr = dispatch_sparse(backend.qr)
 unfold = dispatch_sparse(base.unfold)