Wrap trace for dask

Author: asmeurer

array_api_compat/dask/array/linalg.py

@@ -5,7 +5,7 @@
 
 # Exports
 from dask.array.linalg import * # noqa: F403
-from dask.array import trace, outer
+from dask.array import outer
 
 # These functions are in both the main and linalg namespaces
 from dask.array import matmul, tensordot
@@ -42,6 +42,7 @@ def qr(x: Array, mode: Literal['reduced', 'complete'] = 'reduced',
     if mode != "reduced":
         raise ValueError("dask arrays only support using mode='reduced'")
     return QRResult(*da.linalg.qr(x, **kwargs))
+trace = get_xp(da)(_linalg.trace)
 cholesky = get_xp(da)(_linalg.cholesky)
 matrix_rank = get_xp(da)(_linalg.matrix_rank)
 matrix_norm = get_xp(da)(_linalg.matrix_norm)

dask-xfails.txt

@@ -81,8 +81,6 @@ array_api_tests/test_linalg.py::test_svdvals
 array_api_tests/test_linalg.py::test_cholesky
 # dtype mismatch got uint64, but should be uint8, NPY_PROMOTION_STATE=weak doesn't help :(
 array_api_tests/test_linalg.py::test_tensordot
-# probably same reason for failing as numpy
-array_api_tests/test_linalg.py::test_trace
 
 # AssertionError: out.dtype=uint64, but should be uint8 [tensordot(uint8, uint8)]
 array_api_tests/test_linalg.py::test_linalg_tensordot