Source code for dpctl.tensor._sorting

#                       Data Parallel Control (dpctl)
#
#  Copyright 2020-2024 Intel Corporation
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

import dpctl.tensor as dpt
import dpctl.tensor._tensor_impl as ti
import dpctl.utils as du

from ._numpy_helper import normalize_axis_index
from ._tensor_sorting_impl import (
    _argsort_ascending,
    _argsort_descending,
    _sort_ascending,
    _sort_descending,
)

__all__ = ["sort", "argsort"]


[docs]def sort(x, /, *, axis=-1, descending=False, stable=True): """sort(x, axis=-1, descending=False, stable=True) Returns a sorted copy of an input array `x`. Args: x (usm_ndarray): input array. axis (Optional[int]): axis along which to sort. If set to `-1`, the function must sort along the last axis. Default: `-1`. descending (Optional[bool]): sort order. If `True`, the array must be sorted in descending order (by value). If `False`, the array must be sorted in ascending order (by value). Default: `False`. stable (Optional[bool]): sort stability. If `True`, the returned array must maintain the relative order of `x` values which compare as equal. If `False`, the returned array may or may not maintain the relative order of `x` values which compare as equal. Default: `True`. Returns: usm_ndarray: a sorted array. The returned array has the same data type and the same shape as the input array `x`. """ if not isinstance(x, dpt.usm_ndarray): raise TypeError( f"Expected type dpctl.tensor.usm_ndarray, got {type(x)}" ) nd = x.ndim if nd == 0: axis = normalize_axis_index(axis, ndim=1, msg_prefix="axis") return dpt.copy(x, order="C") else: axis = normalize_axis_index(axis, ndim=nd, msg_prefix="axis") a1 = axis + 1 if a1 == nd: perm = list(range(nd)) arr = x else: perm = [i for i in range(nd) if i != axis] + [ axis, ] arr = dpt.permute_dims(x, perm) exec_q = x.sycl_queue _manager = du.SequentialOrderManager[exec_q] dep_evs = _manager.submitted_events impl_fn = _sort_descending if descending else _sort_ascending if arr.flags.c_contiguous: res = dpt.empty_like(arr, order="C") ht_ev, impl_ev = impl_fn( src=arr, trailing_dims_to_sort=1, dst=res, sycl_queue=exec_q, depends=dep_evs, ) _manager.add_event_pair(ht_ev, impl_ev) else: tmp = dpt.empty_like(arr, order="C") ht_ev, copy_ev = ti._copy_usm_ndarray_into_usm_ndarray( src=arr, dst=tmp, sycl_queue=exec_q, depends=dep_evs ) _manager.add_event_pair(ht_ev, copy_ev) res = dpt.empty_like(arr, order="C") ht_ev, impl_ev = impl_fn( src=tmp, trailing_dims_to_sort=1, dst=res, sycl_queue=exec_q, depends=[copy_ev], ) _manager.add_event_pair(ht_ev, impl_ev) if a1 != nd: inv_perm = sorted(range(nd), key=lambda d: perm[d]) res = dpt.permute_dims(res, inv_perm) return res
[docs]def argsort(x, axis=-1, descending=False, stable=True): """argsort(x, axis=-1, descending=False, stable=True) Returns the indices that sort an array `x` along a specified axis. Args: x (usm_ndarray): input array. axis (Optional[int]): axis along which to sort. If set to `-1`, the function must sort along the last axis. Default: `-1`. descending (Optional[bool]): sort order. If `True`, the array must be sorted in descending order (by value). If `False`, the array must be sorted in ascending order (by value). Default: `False`. stable (Optional[bool]): sort stability. If `True`, the returned array must maintain the relative order of `x` values which compare as equal. If `False`, the returned array may or may not maintain the relative order of `x` values which compare as equal. Default: `True`. Returns: usm_ndarray: an array of indices. The returned array has the same shape as the input array `x`. The return array has default array index data type. """ if not isinstance(x, dpt.usm_ndarray): raise TypeError( f"Expected type dpctl.tensor.usm_ndarray, got {type(x)}" ) nd = x.ndim if nd == 0: axis = normalize_axis_index(axis, ndim=1, msg_prefix="axis") return dpt.zeros_like( x, dtype=ti.default_device_index_type(x.sycl_queue), order="C" ) else: axis = normalize_axis_index(axis, ndim=nd, msg_prefix="axis") a1 = axis + 1 if a1 == nd: perm = list(range(nd)) arr = x else: perm = [i for i in range(nd) if i != axis] + [ axis, ] arr = dpt.permute_dims(x, perm) exec_q = x.sycl_queue _manager = du.SequentialOrderManager[exec_q] dep_evs = _manager.submitted_events impl_fn = _argsort_descending if descending else _argsort_ascending index_dt = ti.default_device_index_type(exec_q) if arr.flags.c_contiguous: res = dpt.empty_like(arr, dtype=index_dt, order="C") ht_ev, impl_ev = impl_fn( src=arr, trailing_dims_to_sort=1, dst=res, sycl_queue=exec_q, depends=dep_evs, ) _manager.add_event_pair(ht_ev, impl_ev) else: tmp = dpt.empty_like(arr, order="C") ht_ev, copy_ev = ti._copy_usm_ndarray_into_usm_ndarray( src=arr, dst=tmp, sycl_queue=exec_q, depends=dep_evs ) _manager.add_event_pair(ht_ev, copy_ev) res = dpt.empty_like(arr, dtype=index_dt, order="C") ht_ev, impl_ev = impl_fn( src=tmp, trailing_dims_to_sort=1, dst=res, sycl_queue=exec_q, depends=[copy_ev], ) _manager.add_event_pair(ht_ev, impl_ev) if a1 != nd: inv_perm = sorted(range(nd), key=lambda d: perm[d]) res = dpt.permute_dims(res, inv_perm) return res