fix argsort dtype

Author: crusaderky

array_api_compat/dask/array/_aliases.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Literal
+from typing import Callable
 
 from ...common import _aliases, array_namespace
 
@@ -31,24 +31,27 @@
 )
 
 from typing import TYPE_CHECKING
+
 if TYPE_CHECKING:
     from typing import Optional, Union
 
-    from ...common._typing import Device, Dtype, Array, NestedSequence, SupportsBufferProtocol
+    from ...common._typing import (
+        Device,
+        Dtype,
+        Array,
+        NestedSequence,
+        SupportsBufferProtocol,
+    )
 
 import dask.array as da
 
 isdtype = get_xp(np)(_aliases.isdtype)
 unstack = get_xp(da)(_aliases.unstack)
 
+
 # da.astype doesn't respect copy=True
 def astype(
-    x: Array,
-    dtype: Dtype,
-    /,
-    *,
-    copy: bool = True,
-    device: Optional[Device] = None
+    x: Array, dtype: Dtype, /, *, copy: bool = True, device: Optional[Device] = None
 ) -> Array:
     """
     Array API compatibility wrapper for astype().
@@ -63,8 +66,10 @@ def astype(
     x = x.astype(dtype)
     return x.copy() if copy else x
 
+
 # Common aliases
 
+
 # This arange func is modified from the common one to
 # not pass stop/step as keyword arguments, which will cause
 # an error with dask
@@ -191,6 +196,7 @@ def asarray(
     concatenate as concat,
 )
 
+
 # dask.array.clip does not work unless all three arguments are provided.
 # Furthermore, the masking workaround in common._aliases.clip cannot work with
 # dask (meaning uint64 promoting to float64 is going to just be unfixed for
@@ -207,8 +213,10 @@ def clip(
     See the corresponding documentation in the array library and/or the array API
     specification for more details.
     """
+
     def _isscalar(a):
         return isinstance(a, (int, float, type(None)))
+
     min_shape = () if _isscalar(min) else min.shape
     max_shape = () if _isscalar(max) else max.shape
 
@@ -231,6 +239,35 @@ def _isscalar(a):
     return astype(da.minimum(da.maximum(x, min), max), x.dtype)
 
 
+def _ensure_single_chunk(x: Array, axis: int) -> tuple[Array, Callable[[Array], Array]]:
+    """
+    Make sure that Array is not broken into multiple chunks along axis.
+
+    Returns
+    -------
+    x : Array
+        The input Array with a single chunk along axis.
+    restore : Callable[Array, Array]
+        function to apply to the output to rechunk it back into reasonable chunks
+    """
+    if axis < 0:
+        axis += x.ndim
+    if x.numblocks[axis] < 2:
+        return x, lambda x: x
+
+    # Break chunks on other axes in an attempt to keep chunk size low
+    x = x.rechunk({i: -1 if i == axis else "auto" for i in range(x.ndim)})
+
+    # Rather than reconstructing the original chunks, which can be a
+    # very expensive affair, just break down oversized chunks without
+    # incurring in any transfers over the network.
+    # This has the downside of a risk of overchunking if the array is
+    # then used in operations against other arrays that match the
+    # original chunking pattern.
+    restore = lambda x: x.rechunk()
+    return x, restore
+
+
 def sort(
     x: Array, /, *, axis: int = -1, descending: bool = False, stable: bool = True
 ) -> Array:
@@ -245,7 +282,20 @@ def sort(
     See the corresponding documentation in the array library and/or the array API
     specification for more details.
     """
-    return _sort_argsort("sort", x, axis=axis, descending=descending, stable=stable)
+    x, restore = _ensure_single_chunk(x, axis)
+
+    meta_xp = array_namespace(x._meta)
+    x = da.map_blocks(
+        meta_xp.sort,
+        x,
+        axis=axis,
+        meta=x._meta,
+        dtype=x.dtype,
+        descending=descending,
+        stable=stable,
+    )
+
+    return restore(x)
 
 
 def argsort(
@@ -262,49 +312,22 @@ def argsort(
     This function temporarily rechunks the array along `axis` into a single chunk.
     This can be extremely inefficient and can lead to out-of-memory errors.
     """
-    return _sort_argsort("argsort", x, axis=axis, descending=descending, stable=stable)
-
-
-def _sort_argsort(
-    func: Literal["sort", "argsort"],
-    x: Array,
-    /,
-    *,
-    axis: int,
-    descending: bool,
-    stable: bool,
-) -> Array:
-    """
-    Implementation of sort() and argsort()
+    x, restore = _ensure_single_chunk(x, axis)
 
-    TODO Implement sort and argsort properly in Dask on top of the shuffle subsystem.
-    """
-    if axis < 0:
-        axis += x.ndim
-    rechunk = False
-    if x.numblocks[axis] > 1:
-        rechunk = True
-        # Break chunks on other axes in an attempt to keep chunk size low
-        x = x.rechunk({i: -1 if i == axis else "auto" for i in range(x.ndim)})
     meta_xp = array_namespace(x._meta)
+    dtype = meta_xp.argsort(x._meta).dtype
+    meta = meta_xp.astype(x._meta, dtype)
     x = da.map_blocks(
-        getattr(meta_xp, func),
+        meta_xp.argsort,
         x,
         axis=axis,
+        meta=meta,
+        dtype=dtype,
         descending=descending,
         stable=stable,
-        dtype=x.dtype,
-        meta=x._meta,
     )
-    if rechunk:
-        # rather than reconstructing the original chunks, which can be a
-        # very expensive affair, just break down oversized chunks without
-        # incurring in any transfers over the network.
-        # This has the downside of a risk of overchunking if the array is
-        # then used in operations against other arrays that match the
-        # original chunking pattern.
-        x = x.rechunk()
-    return x
+
+    return restore(x)
 
 
 _common_aliases = _aliases.__all__
@@ -318,4 +341,4 @@ def _sort_argsort(
                     'complex64', 'complex128', 'iinfo', 'finfo',
                     'can_cast', 'result_type']
 
-_all_ignore = ["Literal", "array_namespace", "get_xp", "da", "np"]
+_all_ignore = ["Callable", "array_namespace", "get_xp", "da", "np"]