numba-cuda 0.22.0__cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl

numba_cuda/numba/cuda/np/numpy_support.py

@@ -0,0 +1,798 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import collections
+import ctypes
+import re
+import numpy as np
+
+from numba.cuda import types
+from numba.cuda.core import errors
+from numba.cuda.typing.templates import signature
+from numba.cuda.np import npdatetime_helpers
+from numba.cuda.core.errors import TypingError
+
+# re-export
+from numba.cuda.cgutils import is_nonelike  # noqa: F401
+
+numpy_version = tuple(map(int, np.__version__.split(".")[:2]))
+
+FROM_DTYPE = {
+    np.dtype("bool"): types.boolean,
+    np.dtype("int8"): types.int8,
+    np.dtype("int16"): types.int16,
+    np.dtype("int32"): types.int32,
+    np.dtype("int64"): types.int64,
+    np.dtype("uint8"): types.uint8,
+    np.dtype("uint16"): types.uint16,
+    np.dtype("uint32"): types.uint32,
+    np.dtype("uint64"): types.uint64,
+    np.dtype("float32"): types.float32,
+    np.dtype("float64"): types.float64,
+    np.dtype("float16"): types.float16,
+    np.dtype("complex64"): types.complex64,
+    np.dtype("complex128"): types.complex128,
+    np.dtype(object): types.pyobject,
+}
+
+re_typestr = re.compile(r"[<>=\|]([a-z])(\d+)?$", re.I)
+re_datetimestr = re.compile(r"[<>=\|]([mM])8?(\[([a-z]+)\])?$", re.I)
+
+sizeof_unicode_char = np.dtype("U1").itemsize
+
+
+def _from_str_dtype(dtype):
+    m = re_typestr.match(dtype.str)
+    if not m:
+        raise errors.NumbaNotImplementedError(dtype)
+    groups = m.groups()
+    typecode = groups[0]
+    if typecode == "U":
+        # unicode
+        if dtype.byteorder not in "=|":
+            raise errors.NumbaNotImplementedError(
+                "Does not support non-native byteorder"
+            )
+        count = dtype.itemsize // sizeof_unicode_char
+        assert count == int(groups[1]), "Unicode char size mismatch"
+        return types.UnicodeCharSeq(count)
+
+    elif typecode == "S":
+        # char
+        count = dtype.itemsize
+        assert count == int(groups[1]), "Char size mismatch"
+        return types.CharSeq(count)
+
+    else:
+        raise errors.NumbaNotImplementedError(dtype)
+
+
+def _from_datetime_dtype(dtype):
+    m = re_datetimestr.match(dtype.str)
+    if not m:
+        raise errors.NumbaNotImplementedError(dtype)
+    groups = m.groups()
+    typecode = groups[0]
+    unit = groups[2] or ""
+    if typecode == "m":
+        return types.NPTimedelta(unit)
+    elif typecode == "M":
+        return types.NPDatetime(unit)
+    else:
+        raise errors.NumbaNotImplementedError(dtype)
+
+
+def from_dtype(dtype):
+    """
+    Return a Numba Type instance corresponding to the given Numpy *dtype*.
+    NumbaNotImplementedError is raised on unsupported Numpy dtypes.
+    """
+    if type(dtype) is type and issubclass(dtype, np.generic):
+        dtype = np.dtype(dtype)
+    elif getattr(dtype, "fields", None) is not None:
+        return from_struct_dtype(dtype)
+
+    try:
+        return FROM_DTYPE[dtype]
+    except KeyError:
+        pass
+
+    try:
+        char = dtype.char
+    except AttributeError:
+        pass
+    else:
+        if char in "SU":
+            return _from_str_dtype(dtype)
+        if char in "mM":
+            return _from_datetime_dtype(dtype)
+        if char in "V" and dtype.subdtype is not None:
+            subtype = from_dtype(dtype.subdtype[0])
+            return types.NestedArray(subtype, dtype.shape)
+
+    raise errors.NumbaNotImplementedError(dtype)
+
+
+_as_dtype_letters = {
+    types.NPDatetime: "M8",
+    types.NPTimedelta: "m8",
+    types.CharSeq: "S",
+    types.UnicodeCharSeq: "U",
+}
+
+
+def as_dtype(nbtype):
+    """
+    Return a numpy dtype instance corresponding to the given Numba type.
+    NotImplementedError is if no correspondence is known.
+    """
+    nbtype = types.unliteral(nbtype)
+    if isinstance(nbtype, (types.Complex, types.Integer, types.Float)):
+        return np.dtype(str(nbtype))
+    if isinstance(nbtype, (types.Boolean)):
+        return np.dtype("?")
+    if isinstance(nbtype, (types.NPDatetime, types.NPTimedelta)):
+        letter = _as_dtype_letters[type(nbtype)]
+        if nbtype.unit:
+            return np.dtype("%s[%s]" % (letter, nbtype.unit))
+        else:
+            return np.dtype(letter)
+    if isinstance(nbtype, (types.CharSeq, types.UnicodeCharSeq)):
+        letter = _as_dtype_letters[type(nbtype)]
+        return np.dtype("%s%d" % (letter, nbtype.count))
+    if isinstance(nbtype, types.Record):
+        return as_struct_dtype(nbtype)
+    if isinstance(nbtype, types.EnumMember):
+        return as_dtype(nbtype.dtype)
+    if isinstance(nbtype, types.npytypes.DType):
+        return as_dtype(nbtype.dtype)
+    if isinstance(nbtype, types.NumberClass):
+        return as_dtype(nbtype.dtype)
+    if isinstance(nbtype, types.NestedArray):
+        spec = (as_dtype(nbtype.dtype), tuple(nbtype.shape))
+        return np.dtype(spec)
+    if isinstance(nbtype, types.PyObject):
+        return np.dtype(object)
+
+    msg = f"{nbtype} cannot be represented as a NumPy dtype"
+    raise errors.NumbaNotImplementedError(msg)
+
+
+def as_struct_dtype(rec):
+    """Convert Numba Record type to NumPy structured dtype"""
+    assert isinstance(rec, types.Record)
+    names = []
+    formats = []
+    offsets = []
+    titles = []
+    # Fill the fields if they are not a title.
+    for k, t in rec.members:
+        if not rec.is_title(k):
+            names.append(k)
+            formats.append(as_dtype(t))
+            offsets.append(rec.offset(k))
+            titles.append(rec.fields[k].title)
+
+    fields = {
+        "names": names,
+        "formats": formats,
+        "offsets": offsets,
+        "itemsize": rec.size,
+        "titles": titles,
+    }
+    _check_struct_alignment(rec, fields)
+    return np.dtype(fields, align=rec.aligned)
+
+
+def _check_struct_alignment(rec, fields):
+    """Check alignment compatibility with Numpy"""
+    if rec.aligned:
+        for k, dt in zip(fields["names"], fields["formats"]):
+            llvm_align = rec.alignof(k)
+            npy_align = dt.alignment
+            if llvm_align is not None and npy_align != llvm_align:
+                msg = (
+                    "NumPy is using a different alignment ({}) "
+                    "than Numba/LLVM ({}) for {}. "
+                    "This is likely a NumPy bug."
+                )
+                raise ValueError(msg.format(npy_align, llvm_align, dt))
+
+
+def map_arrayscalar_type(val):
+    if isinstance(val, np.generic):
+        # We can't blindly call np.dtype() as it loses information
+        # on some types, e.g. datetime64 and timedelta64.
+        dtype = val.dtype
+    else:
+        try:
+            dtype = np.dtype(type(val))
+        except TypeError:
+            raise errors.NumbaNotImplementedError(
+                "no corresponding numpy dtype for %r" % type(val)
+            )
+    return from_dtype(dtype)
+
+
+def is_array(val):
+    return isinstance(val, np.ndarray)
+
+
+def map_layout(val):
+    if val.flags["C_CONTIGUOUS"]:
+        layout = "C"
+    elif val.flags["F_CONTIGUOUS"]:
+        layout = "F"
+    else:
+        layout = "A"
+    return layout
+
+
+def select_array_wrapper(inputs):
+    """
+    Given the array-compatible input types to an operation (e.g. ufunc),
+    select the appropriate input for wrapping the operation output,
+    according to each input's __array_priority__.
+
+    An index into *inputs* is returned.
+    """
+    max_prio = float("-inf")
+    selected_index = None
+    for index, ty in enumerate(inputs):
+        # Ties are broken by choosing the first winner, as in Numpy
+        if (
+            isinstance(ty, types.ArrayCompatible)
+            and ty.array_priority > max_prio
+        ):
+            selected_index = index
+            max_prio = ty.array_priority
+
+    assert selected_index is not None
+    return selected_index
+
+
+def resolve_output_type(context, inputs, formal_output):
+    """
+    Given the array-compatible input types to an operation (e.g. ufunc),
+    and the operation's formal output type (a types.Array instance),
+    resolve the actual output type using the typing *context*.
+
+    This uses a mechanism compatible with Numpy's __array_priority__ /
+    __array_wrap__.
+    """
+    selected_input = inputs[select_array_wrapper(inputs)]
+    args = selected_input, formal_output
+    sig = context.resolve_function_type("__array_wrap__", args, {})
+    if sig is None:
+        if selected_input.array_priority == types.Array.array_priority:
+            # If it's the same priority as a regular array, assume we
+            # should return the output unchanged.
+            # (we can't define __array_wrap__ explicitly for types.Buffer,
+            #  as that would be inherited by most array-compatible objects)
+            return formal_output
+        raise errors.TypingError("__array_wrap__ failed for %s" % (args,))
+    return sig.return_type
+
+
+def supported_ufunc_loop(ufunc, loop):
+    """Return whether the *loop* for the *ufunc* is supported -in nopython-.
+
+    *loop* should be a UFuncLoopSpec instance, and *ufunc* a numpy ufunc.
+
+    For ufuncs implemented using the ufunc_db, it is supported if the ufunc_db
+    contains a lowering definition for 'loop' in the 'ufunc' entry.
+
+    For other ufuncs, it is type based. The loop will be considered valid if it
+    only contains the following letter types: '?bBhHiIlLqQfd'. Note this is
+    legacy and when implementing new ufuncs the ufunc_db should be preferred,
+    as it allows for a more fine-grained incremental support.
+    """
+    # NOTE: Assuming ufunc for the CPUContext
+    from numba.cuda.np import ufunc_db
+
+    loop_sig = loop.ufunc_sig
+    try:
+        # check if the loop has a codegen description in the
+        # ufunc_db. If so, we can proceed.
+
+        # note that as of now not all ufuncs have an entry in the
+        # ufunc_db
+        supported_loop = loop_sig in ufunc_db.get_ufunc_info(ufunc)
+    except KeyError:
+        # for ufuncs not in ufunc_db, base the decision of whether the
+        # loop is supported on its types
+        loop_types = [x.char for x in loop.numpy_inputs + loop.numpy_outputs]
+        supported_types = "?bBhHiIlLqQfd"
+        # check if all the types involved in the ufunc loop are
+        # supported in this mode
+        supported_loop = all(t in supported_types for t in loop_types)
+
+    return supported_loop
+
+
+class UFuncLoopSpec(
+    collections.namedtuple("_UFuncLoopSpec", ("inputs", "outputs", "ufunc_sig"))
+):
+    """
+    An object describing a ufunc loop's inner types.  Properties:
+    - inputs: the inputs' Numba types
+    - outputs: the outputs' Numba types
+    - ufunc_sig: the string representing the ufunc's type signature, in
+      Numpy format (e.g. "ii->i")
+    """
+
+    __slots__ = ()
+
+    @property
+    def numpy_inputs(self):
+        return [as_dtype(x) for x in self.inputs]
+
+    @property
+    def numpy_outputs(self):
+        return [as_dtype(x) for x in self.outputs]
+
+
+def _ufunc_loop_sig(out_tys, in_tys):
+    if len(out_tys) == 1:
+        return signature(out_tys[0], *in_tys)
+    else:
+        return signature(types.Tuple(out_tys), *in_tys)
+
+
+def ufunc_can_cast(from_, to, has_mixed_inputs, casting="safe"):
+    """
+    A variant of np.can_cast() that can allow casting any integer to
+    any real or complex type, in case the operation has mixed-kind
+    inputs.
+
+    For example we want `np.power(float32, int32)` to be computed using
+    SP arithmetic and return `float32`.
+    However, `np.sqrt(int32)` should use DP arithmetic and return `float64`.
+    """
+    from_ = np.dtype(from_)
+    to = np.dtype(to)
+    if has_mixed_inputs and from_.kind in "iu" and to.kind in "cf":
+        # Decide that all integers can cast to any real or complex type.
+        return True
+    return np.can_cast(from_, to, casting)
+
+
+def ufunc_find_matching_loop(ufunc, arg_types):
+    """Find the appropriate loop to be used for a ufunc based on the types
+    of the operands
+
+    ufunc        - The ufunc we want to check
+    arg_types    - The tuple of arguments to the ufunc, including any
+                   explicit output(s).
+    return value - A UFuncLoopSpec identifying the loop, or None
+                   if no matching loop is found.
+    """
+
+    # Separate logical input from explicit output arguments
+    input_types = arg_types[: ufunc.nin]
+    output_types = arg_types[ufunc.nin :]
+    assert len(input_types) == ufunc.nin
+
+    try:
+        np_input_types = [as_dtype(x) for x in input_types]
+    except errors.NumbaNotImplementedError:
+        return None
+    try:
+        np_output_types = [as_dtype(x) for x in output_types]
+    except errors.NumbaNotImplementedError:
+        return None
+
+    # Whether the inputs are mixed integer / floating-point
+    has_mixed_inputs = any(dt.kind in "iu" for dt in np_input_types) and any(
+        dt.kind in "cf" for dt in np_input_types
+    )
+
+    def choose_types(numba_types, ufunc_letters):
+        """
+        Return a list of Numba types representing *ufunc_letters*,
+        except when the letter designates a datetime64 or timedelta64,
+        in which case the type is taken from *numba_types*.
+        """
+        assert len(ufunc_letters) >= len(numba_types)
+        types = [
+            tp if letter in "mM" else from_dtype(np.dtype(letter))
+            for tp, letter in zip(numba_types, ufunc_letters)
+        ]
+        # Add missing types (presumably implicit outputs)
+        types += [
+            from_dtype(np.dtype(letter))
+            for letter in ufunc_letters[len(numba_types) :]
+        ]
+        return types
+
+    def set_output_dt_units(inputs, outputs, ufunc_inputs, ufunc_name):
+        """
+        Sets the output unit of a datetime type based on the input units
+
+        Timedelta is a special dtype that requires the time unit to be
+        specified (day, month, etc). Not every operation with timedelta inputs
+        leads to an output of timedelta output. However, for those that do,
+        the unit of output must be inferred based on the units of the inputs.
+
+        At the moment this function takes care of two cases:
+        a) where all inputs are timedelta with the same unit (mm), and
+        therefore the output has the same unit.
+        This case is used for arr.sum, and for arr1+arr2 where all arrays
+        are timedeltas.
+        If in the future this needs to be extended to a case with mixed units,
+        the rules should be implemented in `npdatetime_helpers` and called
+        from this function to set the correct output unit.
+        b) where left operand is a timedelta, i.e. the "m?" case. This case
+        is used for division, eg timedelta / int.
+
+        At the time of writing, Numba does not support addition of timedelta
+        and other types, so this function does not consider the case "?m",
+        i.e. where timedelta is the right operand to a non-timedelta left
+        operand. To extend it in the future, just add another elif clause.
+        """
+
+        def make_specific(outputs, unit):
+            new_outputs = []
+            for out in outputs:
+                if isinstance(out, types.NPTimedelta) and out.unit == "":
+                    new_outputs.append(types.NPTimedelta(unit))
+                else:
+                    new_outputs.append(out)
+            return new_outputs
+
+        def make_datetime_specific(outputs, dt_unit, td_unit):
+            new_outputs = []
+            for out in outputs:
+                if isinstance(out, types.NPDatetime) and out.unit == "":
+                    unit = npdatetime_helpers.combine_datetime_timedelta_units(
+                        dt_unit, td_unit
+                    )
+                    if unit is None:
+                        raise TypingError(
+                            f"ufunc '{ufunc_name}' is not "
+                            + "supported between "
+                            + f"datetime64[{dt_unit}] "
+                            + f"and timedelta64[{td_unit}]"
+                        )
+                    new_outputs.append(types.NPDatetime(unit))
+                else:
+                    new_outputs.append(out)
+            return new_outputs
+
+        if ufunc_inputs == "mm":
+            if all(inp.unit == inputs[0].unit for inp in inputs):
+                # Case with operation on same units. Operations on different
+                # units not adjusted for now but might need to be
+                # added in the future
+                unit = inputs[0].unit
+                new_outputs = make_specific(outputs, unit)
+            else:
+                return outputs
+            return new_outputs
+        elif ufunc_inputs == "mM":
+            # case where the left operand has timedelta type
+            # and the right operand has datetime
+            td_unit = inputs[0].unit
+            dt_unit = inputs[1].unit
+            return make_datetime_specific(outputs, dt_unit, td_unit)
+
+        elif ufunc_inputs == "Mm":
+            # case where the right operand has timedelta type
+            # and the left operand has datetime
+            dt_unit = inputs[0].unit
+            td_unit = inputs[1].unit
+            return make_datetime_specific(outputs, dt_unit, td_unit)
+
+        elif ufunc_inputs[0] == "m":
+            # case where the left operand has timedelta type
+            unit = inputs[0].unit
+            new_outputs = make_specific(outputs, unit)
+            return new_outputs
+
+    # In NumPy, the loops are evaluated from first to last. The first one
+    # that is viable is the one used. One loop is viable if it is possible
+    # to cast every input operand to the one expected by the ufunc.
+    # Also under NumPy 1.10+ the output must be able to be cast back
+    # to a close enough type ("same_kind").
+
+    for candidate in ufunc.types:
+        ufunc_inputs = candidate[: ufunc.nin]
+        ufunc_outputs = candidate[-ufunc.nout :] if ufunc.nout else []
+
+        if "e" in ufunc_inputs:
+            # Skip float16 arrays since we don't have implementation for them
+            continue
+        if "O" in ufunc_inputs:
+            # Skip object arrays
+            continue
+        found = True
+        # Skip if any input or output argument is mismatching
+        for outer, inner in zip(np_input_types, ufunc_inputs):
+            # (outer is a dtype instance, inner is a type char)
+            if outer.char in "mM" or inner in "mM":
+                # For datetime64 and timedelta64, we want to retain
+                # precise typing (i.e. the units); therefore we look for
+                # an exact match.
+                if outer.char != inner:
+                    found = False
+                    break
+            elif not ufunc_can_cast(
+                outer.char, inner, has_mixed_inputs, "safe"
+            ):
+                found = False
+                break
+        if found:
+            # Can we cast the inner result to the outer result type?
+            for outer, inner in zip(np_output_types, ufunc_outputs):
+                if outer.char not in "mM" and not ufunc_can_cast(
+                    inner, outer.char, has_mixed_inputs, "same_kind"
+                ):
+                    found = False
+                    break
+        if found:
+            # Found: determine the Numba types for the loop's inputs and
+            # outputs.
+            try:
+                inputs = choose_types(input_types, ufunc_inputs)
+                outputs = choose_types(output_types, ufunc_outputs)
+                # if the left operand or both are timedeltas, or the first
+                # argument is datetime and the second argument is timedelta,
+                # then the output units need to be determined.
+                if ufunc_inputs[0] == "m" or ufunc_inputs == "Mm":
+                    outputs = set_output_dt_units(
+                        inputs, outputs, ufunc_inputs, ufunc.__name__
+                    )
+
+            except errors.NumbaNotImplementedError:
+                # One of the selected dtypes isn't supported by Numba
+                # (e.g. float16), try other candidates
+                continue
+            else:
+                return UFuncLoopSpec(inputs, outputs, candidate)
+
+    return None
+
+
+def _is_aligned_struct(struct):
+    return struct.isalignedstruct
+
+
+def from_struct_dtype(dtype):
+    """Convert a NumPy structured dtype to Numba Record type"""
+    if dtype.hasobject:
+        msg = "dtypes that contain object are not supported."
+        raise errors.NumbaNotImplementedError(msg)
+
+    fields = []
+    for name, info in dtype.fields.items():
+        # *info* may have 3 element
+        [elemdtype, offset] = info[:2]
+        title = info[2] if len(info) == 3 else None
+
+        ty = from_dtype(elemdtype)
+        infos = {
+            "type": ty,
+            "offset": offset,
+            "title": title,
+        }
+        fields.append((name, infos))
+
+    # Note: dtype.alignment is not consistent.
+    #       It is different after passing into a recarray.
+    #       recarray(N, dtype=mydtype).dtype.alignment != mydtype.alignment
+    size = dtype.itemsize
+    aligned = _is_aligned_struct(dtype)
+
+    return types.Record(fields, size, aligned)
+
+
+def _get_bytes_buffer(ptr, nbytes):
+    """
+    Get a ctypes array of *nbytes* starting at *ptr*.
+    """
+    if isinstance(ptr, ctypes.c_void_p):
+        ptr = ptr.value
+    arrty = ctypes.c_byte * nbytes
+    return arrty.from_address(ptr)
+
+
+def _get_array_from_ptr(ptr, nbytes, dtype):
+    return np.frombuffer(_get_bytes_buffer(ptr, nbytes), dtype)
+
+
+def carray(ptr, shape, dtype=None):
+    """
+    Return a Numpy array view over the data pointed to by *ptr* with the
+    given *shape*, in C order.  If *dtype* is given, it is used as the
+    array's dtype, otherwise the array's dtype is inferred from *ptr*'s type.
+    """
+    from numba.cuda.typing.ctypes_utils import from_ctypes
+
+    try:
+        # Use ctypes parameter protocol if available
+        ptr = ptr._as_parameter_
+    except AttributeError:
+        pass
+
+    # Normalize dtype, to accept e.g. "int64" or np.int64
+    if dtype is not None:
+        dtype = np.dtype(dtype)
+
+    if isinstance(ptr, ctypes.c_void_p):
+        if dtype is None:
+            raise TypeError("explicit dtype required for void* argument")
+        p = ptr
+    elif isinstance(ptr, ctypes._Pointer):
+        ptrty = from_ctypes(ptr.__class__)
+        assert isinstance(ptrty, types.CPointer)
+        ptr_dtype = as_dtype(ptrty.dtype)
+        if dtype is not None and dtype != ptr_dtype:
+            raise TypeError(
+                "mismatching dtype '%s' for pointer %s" % (dtype, ptr)
+            )
+        dtype = ptr_dtype
+        p = ctypes.cast(ptr, ctypes.c_void_p)
+    else:
+        raise TypeError("expected a ctypes pointer, got %r" % (ptr,))
+
+    nbytes = dtype.itemsize * np.prod(shape, dtype=np.intp)
+    return _get_array_from_ptr(p, nbytes, dtype).reshape(shape)
+
+
+def farray(ptr, shape, dtype=None):
+    """
+    Return a Numpy array view over the data pointed to by *ptr* with the
+    given *shape*, in Fortran order.  If *dtype* is given, it is used as the
+    array's dtype, otherwise the array's dtype is inferred from *ptr*'s type.
+    """
+    if not isinstance(shape, int):
+        shape = shape[::-1]
+    return carray(ptr, shape, dtype).T
+
+
+def is_contiguous(dims, strides, itemsize):
+    """Is the given shape, strides, and itemsize of C layout?
+
+    Note: The code is usable as a numba-compiled function
+    """
+    nd = len(dims)
+    # Check and skip 1s or 0s in inner dims
+    innerax = nd - 1
+    while innerax > -1 and dims[innerax] <= 1:
+        innerax -= 1
+
+    # Early exit if all axis are 1s or 0s
+    if innerax < 0:
+        return True
+
+    # Check itemsize matches innermost stride
+    if itemsize != strides[innerax]:
+        return False
+
+    # Check and skip 1s or 0s in outer dims
+    outerax = 0
+    while outerax < innerax and dims[outerax] <= 1:
+        outerax += 1
+
+    # Check remaining strides to be contiguous
+    ax = innerax
+    while ax > outerax:
+        if strides[ax] * dims[ax] != strides[ax - 1]:
+            return False
+        ax -= 1
+    return True
+
+
+def is_fortran(dims, strides, itemsize):
+    """Is the given shape, strides, and itemsize of F layout?
+
+    Note: The code is usable as a numba-compiled function
+    """
+    nd = len(dims)
+    # Check and skip 1s or 0s in inner dims
+    firstax = 0
+    while firstax < nd and dims[firstax] <= 1:
+        firstax += 1
+
+    # Early exit if all axis are 1s or 0s
+    if firstax >= nd:
+        return True
+
+    # Check itemsize matches innermost stride
+    if itemsize != strides[firstax]:
+        return False
+
+    # Check and skip 1s or 0s in outer dims
+    lastax = nd - 1
+    while lastax > firstax and dims[lastax] <= 1:
+        lastax -= 1
+
+    # Check remaining strides to be contiguous
+    ax = firstax
+    while ax < lastax:
+        if strides[ax] * dims[ax] != strides[ax + 1]:
+            return False
+        ax += 1
+    return True
+
+
+def type_can_asarray(arr):
+    """Returns True if the type of 'arr' is supported by the Numba `np.asarray`
+    implementation, False otherwise.
+    """
+
+    ok = (
+        types.Array,
+        types.Sequence,
+        types.Tuple,
+        types.StringLiteral,
+        types.Number,
+        types.Boolean,
+        types.containers.ListType,
+    )
+
+    return isinstance(arr, ok)
+
+
+def type_is_scalar(typ):
+    """Returns True if the type of 'typ' is a scalar type, according to
+    NumPy rules. False otherwise.
+    https://numpy.org/doc/stable/reference/arrays.scalars.html#built-in-scalar-types
+    """
+
+    ok = (
+        types.Boolean,
+        types.Number,
+        types.UnicodeType,
+        types.StringLiteral,
+        types.NPTimedelta,
+        types.NPDatetime,
+    )
+    return isinstance(typ, ok)
+
+
+def check_is_integer(v, name):
+    """Raises TypingError if the value is not an integer."""
+    if not isinstance(v, (int, types.Integer)):
+        raise TypingError("{} must be an integer".format(name))
+
+
+def lt_floats(a, b):
+    # Adapted from NumPy commit 717c7acf which introduced the behavior of
+    # putting NaNs at the end.
+    # The code is later moved to numpy/core/src/npysort/npysort_common.h
+    # This info is gathered as of NumPy commit d8c09c50
+    return a < b or (np.isnan(b) and not np.isnan(a))
+
+
+def lt_complex(a, b):
+    if np.isnan(a.real):
+        if np.isnan(b.real):
+            if np.isnan(a.imag):
+                return False
+            else:
+                if np.isnan(b.imag):
+                    return True
+                else:
+                    return a.imag < b.imag
+        else:
+            return False
+
+    else:
+        if np.isnan(b.real):
+            return True
+        else:
+            if np.isnan(a.imag):
+                if np.isnan(b.imag):
+                    return a.real < b.real
+                else:
+                    return False
+            else:
+                if np.isnan(b.imag):
+                    return True
+                else:
+                    if a.real < b.real:
+                        return True
+                    elif a.real == b.real:
+                        return a.imag < b.imag
+                    return False