wrap svd

Author: lithomas1

array_api_compat/dask/array/_aliases.py

@@ -12,7 +12,7 @@
 if TYPE_CHECKING:
     from typing import Optional, Tuple, Union
 
-    from ...common._typing import Device, Dtype, ndarray
+    from ...common._typing import Device, Dtype, Array
 
 import dask.array as da
 
@@ -37,7 +37,7 @@ def dask_arange(
     dtype: Optional[Dtype] = None,
     device: Optional[Device] = None,
     **kwargs,
-) -> ndarray:
+) -> Array:
     _check_device(xp, device)
     args = [start]
     if stop is not None:
@@ -99,8 +99,18 @@ def dask_arange(
 matrix_rank = get_xp(da)(_linalg.matrix_rank)
 matrix_norm = get_xp(da)(_linalg.matrix_norm)
 
+# Wrap the svd functions to not pass full_matrices to dask
+# when full_matrices=False (as that is the defualt behavior for dask),
+# and dask doesn't have the full_matrices keyword
+_svd = get_xp(da)(_linalg.svd)
 
-def svdvals(x: ndarray) -> Union[ndarray, Tuple[ndarray, ...]]:
+def svd(x: Array, full_matrices: bool = True, **kwargs) -> SVDResult:
+    if full_matrices:
+        return _svd(x, full_matrices=full_matrices, **kwargs)
+    return _svd(x, **kwargs)
+
+
+def svdvals(x: Array) -> Array:
     # TODO: can't avoid computing U or V for dask
     _, s, _ = da.linalg.svd(x)
     return s

dask-xfails.txt

@@ -112,7 +112,6 @@ array_api_tests/test_linalg.py::test_solve
 # missing full_matrics kw
 # https://github.com/dask/dask/issues/10389
 # also only supports 2-d inputs
-array_api_tests/test_signatures.py::test_extension_func_signature[linalg.svd]
 array_api_tests/test_linalg.py::test_svd
 
 # Missing dlpack stuff