Fix ruff errors for dask/array/linalg

Author: adonath

array_api_compat/dask/array/_aliases.py

@@ -1,17 +1,18 @@
 from __future__ import annotations
+
 from functools import partial
+from typing import TYPE_CHECKING
 
-from ...common import _aliases
-from ...common._helpers import _check_device
+import numpy as np
 
 from ..._internal import get_xp
+from ...common import _aliases, _linalg
+from ...common._helpers import _check_device
 
-import numpy as np
-
-from typing import TYPE_CHECKING
 if TYPE_CHECKING:
-    from typing import Optional, Union
-    from ...common._typing import ndarray, Device, Dtype
+    from typing import Optional, Tuple, Union
+
+    from ...common._typing import Device, Dtype, ndarray
 
 import dask.array as da
 
@@ -24,6 +25,7 @@
 # not pass stop/step as keyword arguments, which will cause
 # an error with dask
 
+
 # TODO: delete the xp stuff, it shouldn't be necessary
 def dask_arange(
     start: Union[int, float],
@@ -34,7 +36,7 @@ def dask_arange(
     xp,
     dtype: Optional[Dtype] = None,
     device: Optional[Device] = None,
-    **kwargs
+    **kwargs,
 ) -> ndarray:
     _check_device(xp, device)
     args = [start]
@@ -47,10 +49,11 @@ def dask_arange(
     args.append(step)
     return xp.arange(*args, dtype=dtype, **kwargs)
 
+
 arange = get_xp(da)(dask_arange)
 eye = get_xp(da)(_aliases.eye)
 
-asarray = partial(_aliases._asarray, namespace='dask.array')
+asarray = partial(_aliases._asarray, namespace="dask.array")
 asarray.__doc__ = _aliases._asarray.__doc__
 
 linspace = get_xp(da)(_aliases.linspace)
@@ -86,3 +89,22 @@ def dask_arange(
 matmul = get_xp(np)(_aliases.matmul)
 tensordot = get_xp(np)(_aliases.tensordot)
 
+
+EighResult = _linalg.EighResult
+QRResult = _linalg.QRResult
+SlogdetResult = _linalg.SlogdetResult
+SVDResult = _linalg.SVDResult
+qr = get_xp(da)(_linalg.qr)
+cholesky = get_xp(da)(_linalg.cholesky)
+matrix_rank = get_xp(da)(_linalg.matrix_rank)
+matrix_norm = get_xp(da)(_linalg.matrix_norm)
+
+
+def svdvals(x: ndarray) -> Union[ndarray, Tuple[ndarray, ...]]:
+    # TODO: can't avoid computing U or V for dask
+    _, s, _ = da.linalg.svd(x)
+    return s
+
+
+vector_norm = get_xp(da)(_linalg.vector_norm)
+diagonal = get_xp(da)(_linalg.diagonal)

array_api_compat/dask/array/linalg.py

@@ -1,48 +1,50 @@
-from __future__ import annotations
-
-from dask.array.linalg import *
-from ...common import _linalg
-from ..._internal import get_xp
-from dask.array import matmul, tensordot, trace, outer
-from ._aliases import matrix_transpose, vecdot
-
-import dask.array as da
-
-from typing import TYPE_CHECKING
-if TYPE_CHECKING:
-    from typing import Union, Tuple
-    from ...common._typing import ndarray
-
-# cupy.linalg doesn't have __all__. If it is added, replace this with
-#
-# from cupy.linalg import __all__ as linalg_all
-_n = {}
-exec('from dask.array.linalg import *', _n)
-del _n['__builtins__']
-linalg_all = list(_n)
-del _n
-
-EighResult = _linalg.EighResult
-QRResult = _linalg.QRResult
-SlogdetResult = _linalg.SlogdetResult
-SVDResult = _linalg.SVDResult
-qr = get_xp(da)(_linalg.qr)
-cholesky = get_xp(da)(_linalg.cholesky)
-matrix_rank = get_xp(da)(_linalg.matrix_rank)
-matrix_norm = get_xp(da)(_linalg.matrix_norm)
-
-def svdvals(x: ndarray) -> Union[ndarray, Tuple[ndarray, ...]]:
-    # TODO: can't avoid computing U or V for dask
-    _, s, _ =  svd(x)
-    return s
-
-vector_norm = get_xp(da)(_linalg.vector_norm)
-diagonal = get_xp(da)(_linalg.diagonal)
-
-__all__ = linalg_all + ["EighResult", "QRResult", "SlogdetResult",
-                        "SVDResult", "qr", "cholesky", "matrix_rank", "matrix_norm",
-                        "svdvals", "vector_norm", "diagonal"]
-
-del get_xp
-del da
-del _linalg
+import dask.array as _da
+from dask.array import (
+    matmul,
+    outer,
+    tensordot,
+    trace,
+)
+from dask.array.linalg import *  # noqa: F401, F403
+
+from .._internal import _get_all_public_members
+from ._aliases import (
+    EighResult,
+    QRResult,
+    SlogdetResult,
+    SVDResult,
+    cholesky,
+    diagonal,
+    matrix_norm,
+    matrix_rank,
+    matrix_transpose,
+    qr,
+    svdvals,
+    vecdot,
+    vector_norm,
+)
+
+__all__ = [
+    "matmul",
+    "outer",
+    "tensordot",
+    "trace",
+]
+
+__all__ += _get_all_public_members(_da.linalg)
+
+__all__ += [
+    "EighResult",
+    "QRResult",
+    "SlogdetResult",
+    "SVDResult",
+    "qr",
+    "cholesky",
+    "matrix_rank",
+    "matrix_norm",
+    "matrix_transpose",
+    "vecdot",
+    "svdvals",
+    "vector_norm",
+    "diagonal",
+]