6866 lines
240 KiB
Python
6866 lines
240 KiB
Python
"""
|
|
Implementation of operations on Array objects and objects supporting
|
|
the buffer protocol.
|
|
"""
|
|
|
|
import functools
|
|
import math
|
|
import operator
|
|
|
|
from llvmlite import ir
|
|
from llvmlite.ir import Constant
|
|
|
|
import numpy as np
|
|
|
|
from numba import pndindex, literal_unroll
|
|
from numba.core import types, typing, errors, cgutils, extending
|
|
from numba.np.numpy_support import (as_dtype, from_dtype, carray, farray,
|
|
is_contiguous, is_fortran,
|
|
check_is_integer, type_is_scalar,
|
|
lt_complex, lt_floats)
|
|
from numba.np.numpy_support import type_can_asarray, is_nonelike, numpy_version
|
|
from numba.core.imputils import (lower_builtin, lower_getattr,
|
|
lower_getattr_generic,
|
|
lower_setattr_generic,
|
|
lower_cast, lower_constant,
|
|
iternext_impl, impl_ret_borrowed,
|
|
impl_ret_new_ref, impl_ret_untracked,
|
|
RefType)
|
|
from numba.core.typing import signature
|
|
from numba.core.types import StringLiteral
|
|
from numba.core.extending import (register_jitable, overload, overload_method,
|
|
intrinsic, overload_attribute)
|
|
from numba.misc import quicksort, mergesort
|
|
from numba.cpython import slicing
|
|
from numba.cpython.unsafe.tuple import tuple_setitem, build_full_slice_tuple
|
|
from numba.core.extending import overload_classmethod
|
|
from numba.core.typing.npydecl import (parse_dtype as ty_parse_dtype,
|
|
parse_shape as ty_parse_shape,
|
|
_parse_nested_sequence,
|
|
_sequence_of_arrays,
|
|
_choose_concatenation_layout)
|
|
|
|
|
|
def set_range_metadata(builder, load, lower_bound, upper_bound):
|
|
"""
|
|
Set the "range" metadata on a load instruction.
|
|
Note the interval is in the form [lower_bound, upper_bound).
|
|
"""
|
|
range_operands = [Constant(load.type, lower_bound),
|
|
Constant(load.type, upper_bound)]
|
|
md = builder.module.add_metadata(range_operands)
|
|
load.set_metadata("range", md)
|
|
|
|
|
|
def mark_positive(builder, load):
|
|
"""
|
|
Mark the result of a load instruction as positive (or zero).
|
|
"""
|
|
upper_bound = (1 << (load.type.width - 1)) - 1
|
|
set_range_metadata(builder, load, 0, upper_bound)
|
|
|
|
|
|
def make_array(array_type):
|
|
"""
|
|
Return the Structure representation of the given *array_type*
|
|
(an instance of types.ArrayCompatible).
|
|
|
|
Note this does not call __array_wrap__ in case a new array structure
|
|
is being created (rather than populated).
|
|
"""
|
|
real_array_type = array_type.as_array
|
|
base = cgutils.create_struct_proxy(real_array_type)
|
|
ndim = real_array_type.ndim
|
|
|
|
class ArrayStruct(base):
|
|
|
|
def _make_refs(self, ref):
|
|
sig = signature(real_array_type, array_type)
|
|
try:
|
|
array_impl = self._context.get_function('__array__', sig)
|
|
except NotImplementedError:
|
|
return super(ArrayStruct, self)._make_refs(ref)
|
|
|
|
# Return a wrapped structure and its unwrapped reference
|
|
datamodel = self._context.data_model_manager[array_type]
|
|
be_type = self._get_be_type(datamodel)
|
|
if ref is None:
|
|
outer_ref = cgutils.alloca_once(self._builder, be_type,
|
|
zfill=True)
|
|
else:
|
|
outer_ref = ref
|
|
# NOTE: __array__ is called with a pointer and expects a pointer
|
|
# in return!
|
|
ref = array_impl(self._builder, (outer_ref,))
|
|
return outer_ref, ref
|
|
|
|
@property
|
|
def shape(self):
|
|
"""
|
|
Override .shape to inform LLVM that its elements are all positive.
|
|
"""
|
|
builder = self._builder
|
|
if ndim == 0:
|
|
return base.__getattr__(self, "shape")
|
|
|
|
# Unfortunately, we can't use llvm.assume as its presence can
|
|
# seriously pessimize performance,
|
|
# *and* the range metadata currently isn't improving anything here,
|
|
# see https://llvm.org/bugs/show_bug.cgi?id=23848 !
|
|
ptr = self._get_ptr_by_name("shape")
|
|
dims = []
|
|
for i in range(ndim):
|
|
dimptr = cgutils.gep_inbounds(builder, ptr, 0, i)
|
|
load = builder.load(dimptr)
|
|
dims.append(load)
|
|
mark_positive(builder, load)
|
|
|
|
return cgutils.pack_array(builder, dims)
|
|
|
|
return ArrayStruct
|
|
|
|
|
|
def get_itemsize(context, array_type):
|
|
"""
|
|
Return the item size for the given array or buffer type.
|
|
"""
|
|
llty = context.get_data_type(array_type.dtype)
|
|
return context.get_abi_sizeof(llty)
|
|
|
|
|
|
def load_item(context, builder, arrayty, ptr):
|
|
"""
|
|
Load the item at the given array pointer.
|
|
"""
|
|
align = None if arrayty.aligned else 1
|
|
return context.unpack_value(builder, arrayty.dtype, ptr,
|
|
align=align)
|
|
|
|
|
|
def store_item(context, builder, arrayty, val, ptr):
|
|
"""
|
|
Store the item at the given array pointer.
|
|
"""
|
|
align = None if arrayty.aligned else 1
|
|
return context.pack_value(builder, arrayty.dtype, val, ptr, align=align)
|
|
|
|
|
|
def fix_integer_index(context, builder, idxty, idx, size):
|
|
"""
|
|
Fix the integer index' type and value for the given dimension size.
|
|
"""
|
|
if idxty.signed:
|
|
ind = context.cast(builder, idx, idxty, types.intp)
|
|
ind = slicing.fix_index(builder, ind, size)
|
|
else:
|
|
ind = context.cast(builder, idx, idxty, types.uintp)
|
|
return ind
|
|
|
|
|
|
def normalize_index(context, builder, idxty, idx):
|
|
"""
|
|
Normalize the index type and value. 0-d arrays are converted to scalars.
|
|
"""
|
|
if isinstance(idxty, types.Array) and idxty.ndim == 0:
|
|
assert isinstance(idxty.dtype, types.Integer)
|
|
idxary = make_array(idxty)(context, builder, idx)
|
|
idxval = load_item(context, builder, idxty, idxary.data)
|
|
return idxty.dtype, idxval
|
|
else:
|
|
return idxty, idx
|
|
|
|
|
|
def normalize_indices(context, builder, index_types, indices):
|
|
"""
|
|
Same as normalize_index(), but operating on sequences of
|
|
index types and values.
|
|
"""
|
|
if len(indices):
|
|
index_types, indices = zip(*[normalize_index(context, builder, idxty,
|
|
idx)
|
|
for idxty, idx in zip(index_types, indices)
|
|
])
|
|
return index_types, indices
|
|
|
|
|
|
def populate_array(array, data, shape, strides, itemsize, meminfo,
|
|
parent=None):
|
|
"""
|
|
Helper function for populating array structures.
|
|
This avoids forgetting to set fields.
|
|
|
|
*shape* and *strides* can be Python tuples or LLVM arrays.
|
|
"""
|
|
context = array._context
|
|
builder = array._builder
|
|
datamodel = array._datamodel
|
|
# doesn't matter what this array type instance is, it's just to get the
|
|
# fields for the datamodel of the standard array type in this context
|
|
standard_array = types.Array(types.float64, 1, 'C')
|
|
standard_array_type_datamodel = context.data_model_manager[standard_array]
|
|
required_fields = set(standard_array_type_datamodel._fields)
|
|
datamodel_fields = set(datamodel._fields)
|
|
# Make sure that the presented array object has a data model that is close
|
|
# enough to an array for this function to proceed.
|
|
if (required_fields & datamodel_fields) != required_fields:
|
|
missing = required_fields - datamodel_fields
|
|
msg = (f"The datamodel for type {array._fe_type} is missing "
|
|
f"field{'s' if len(missing) > 1 else ''} {missing}.")
|
|
raise ValueError(msg)
|
|
|
|
if meminfo is None:
|
|
meminfo = Constant(context.get_value_type(
|
|
datamodel.get_type('meminfo')), None)
|
|
|
|
intp_t = context.get_value_type(types.intp)
|
|
if isinstance(shape, (tuple, list)):
|
|
shape = cgutils.pack_array(builder, shape, intp_t)
|
|
if isinstance(strides, (tuple, list)):
|
|
strides = cgutils.pack_array(builder, strides, intp_t)
|
|
if isinstance(itemsize, int):
|
|
itemsize = intp_t(itemsize)
|
|
|
|
attrs = dict(shape=shape,
|
|
strides=strides,
|
|
data=data,
|
|
itemsize=itemsize,
|
|
meminfo=meminfo,)
|
|
|
|
# Set `parent` attribute
|
|
if parent is None:
|
|
attrs['parent'] = Constant(context.get_value_type(
|
|
datamodel.get_type('parent')), None)
|
|
else:
|
|
attrs['parent'] = parent
|
|
# Calc num of items from shape
|
|
nitems = context.get_constant(types.intp, 1)
|
|
unpacked_shape = cgutils.unpack_tuple(builder, shape, shape.type.count)
|
|
# (note empty shape => 0d array therefore nitems = 1)
|
|
for axlen in unpacked_shape:
|
|
nitems = builder.mul(nitems, axlen, flags=['nsw'])
|
|
attrs['nitems'] = nitems
|
|
|
|
# Make sure that we have all the fields
|
|
got_fields = set(attrs.keys())
|
|
if got_fields != required_fields:
|
|
raise ValueError("missing {0}".format(required_fields - got_fields))
|
|
|
|
# Set field value
|
|
for k, v in attrs.items():
|
|
setattr(array, k, v)
|
|
|
|
return array
|
|
|
|
|
|
def update_array_info(aryty, array):
|
|
"""
|
|
Update some auxiliary information in *array* after some of its fields
|
|
were changed. `itemsize` and `nitems` are updated.
|
|
"""
|
|
context = array._context
|
|
builder = array._builder
|
|
|
|
# Calc num of items from shape
|
|
nitems = context.get_constant(types.intp, 1)
|
|
unpacked_shape = cgutils.unpack_tuple(builder, array.shape, aryty.ndim)
|
|
for axlen in unpacked_shape:
|
|
nitems = builder.mul(nitems, axlen, flags=['nsw'])
|
|
array.nitems = nitems
|
|
|
|
array.itemsize = context.get_constant(types.intp,
|
|
get_itemsize(context, aryty))
|
|
|
|
|
|
def normalize_axis(func_name, arg_name, ndim, axis):
|
|
"""Constrain axis values to valid positive values."""
|
|
raise NotImplementedError()
|
|
|
|
|
|
@overload(normalize_axis)
|
|
def normalize_axis_overloads(func_name, arg_name, ndim, axis):
|
|
if not isinstance(func_name, StringLiteral):
|
|
raise errors.TypingError("func_name must be a str literal.")
|
|
if not isinstance(arg_name, StringLiteral):
|
|
raise errors.TypingError("arg_name must be a str literal.")
|
|
|
|
msg = (
|
|
f"{func_name.literal_value}: Argument {arg_name.literal_value} "
|
|
"out of bounds for dimensions of the array"
|
|
)
|
|
|
|
def impl(func_name, arg_name, ndim, axis):
|
|
if axis < 0:
|
|
axis += ndim
|
|
if axis < 0 or axis >= ndim:
|
|
raise ValueError(msg)
|
|
|
|
return axis
|
|
|
|
return impl
|
|
|
|
|
|
@lower_builtin('getiter', types.Buffer)
|
|
def getiter_array(context, builder, sig, args):
|
|
[arrayty] = sig.args
|
|
[array] = args
|
|
|
|
iterobj = context.make_helper(builder, sig.return_type)
|
|
|
|
zero = context.get_constant(types.intp, 0)
|
|
indexptr = cgutils.alloca_once_value(builder, zero)
|
|
|
|
iterobj.index = indexptr
|
|
iterobj.array = array
|
|
|
|
# Incref array
|
|
if context.enable_nrt:
|
|
context.nrt.incref(builder, arrayty, array)
|
|
|
|
res = iterobj._getvalue()
|
|
|
|
# Note: a decref on the iterator will dereference all internal MemInfo*
|
|
out = impl_ret_new_ref(context, builder, sig.return_type, res)
|
|
return out
|
|
|
|
|
|
def _getitem_array_single_int(context, builder, return_type, aryty, ary, idx):
|
|
""" Evaluate `ary[idx]`, where idx is a single int. """
|
|
# optimized form of _getitem_array_generic
|
|
shapes = cgutils.unpack_tuple(builder, ary.shape, count=aryty.ndim)
|
|
strides = cgutils.unpack_tuple(builder, ary.strides, count=aryty.ndim)
|
|
offset = builder.mul(strides[0], idx)
|
|
dataptr = cgutils.pointer_add(builder, ary.data, offset)
|
|
view_shapes = shapes[1:]
|
|
view_strides = strides[1:]
|
|
|
|
if isinstance(return_type, types.Buffer):
|
|
# Build array view
|
|
retary = make_view(context, builder, aryty, ary, return_type,
|
|
dataptr, view_shapes, view_strides)
|
|
return retary._getvalue()
|
|
else:
|
|
# Load scalar from 0-d result
|
|
assert not view_shapes
|
|
return load_item(context, builder, aryty, dataptr)
|
|
|
|
|
|
@lower_builtin('iternext', types.ArrayIterator)
|
|
@iternext_impl(RefType.BORROWED)
|
|
def iternext_array(context, builder, sig, args, result):
|
|
[iterty] = sig.args
|
|
[iter] = args
|
|
arrayty = iterty.array_type
|
|
|
|
iterobj = context.make_helper(builder, iterty, value=iter)
|
|
ary = make_array(arrayty)(context, builder, value=iterobj.array)
|
|
|
|
nitems, = cgutils.unpack_tuple(builder, ary.shape, count=1)
|
|
|
|
index = builder.load(iterobj.index)
|
|
is_valid = builder.icmp_signed('<', index, nitems)
|
|
result.set_valid(is_valid)
|
|
|
|
with builder.if_then(is_valid):
|
|
value = _getitem_array_single_int(
|
|
context, builder, iterty.yield_type, arrayty, ary, index
|
|
)
|
|
result.yield_(value)
|
|
nindex = cgutils.increment_index(builder, index)
|
|
builder.store(nindex, iterobj.index)
|
|
|
|
|
|
# ------------------------------------------------------------------------------
|
|
# Basic indexing (with integers and slices only)
|
|
|
|
def basic_indexing(context, builder, aryty, ary, index_types, indices,
|
|
boundscheck=None):
|
|
"""
|
|
Perform basic indexing on the given array.
|
|
A (data pointer, shapes, strides) tuple is returned describing
|
|
the corresponding view.
|
|
"""
|
|
zero = context.get_constant(types.intp, 0)
|
|
one = context.get_constant(types.intp, 1)
|
|
|
|
shapes = cgutils.unpack_tuple(builder, ary.shape, aryty.ndim)
|
|
strides = cgutils.unpack_tuple(builder, ary.strides, aryty.ndim)
|
|
|
|
output_indices = []
|
|
output_shapes = []
|
|
output_strides = []
|
|
|
|
num_newaxes = len([idx for idx in index_types if is_nonelike(idx)])
|
|
ax = 0
|
|
for indexval, idxty in zip(indices, index_types):
|
|
if idxty is types.ellipsis:
|
|
# Fill up missing dimensions at the middle
|
|
n_missing = aryty.ndim - len(indices) + 1 + num_newaxes
|
|
for i in range(n_missing):
|
|
output_indices.append(zero)
|
|
output_shapes.append(shapes[ax])
|
|
output_strides.append(strides[ax])
|
|
ax += 1
|
|
continue
|
|
# Regular index value
|
|
if isinstance(idxty, types.SliceType):
|
|
slice = context.make_helper(builder, idxty, value=indexval)
|
|
slicing.guard_invalid_slice(context, builder, idxty, slice)
|
|
slicing.fix_slice(builder, slice, shapes[ax])
|
|
output_indices.append(slice.start)
|
|
sh = slicing.get_slice_length(builder, slice)
|
|
st = slicing.fix_stride(builder, slice, strides[ax])
|
|
output_shapes.append(sh)
|
|
output_strides.append(st)
|
|
elif isinstance(idxty, types.Integer):
|
|
ind = fix_integer_index(context, builder, idxty, indexval,
|
|
shapes[ax])
|
|
if boundscheck:
|
|
cgutils.do_boundscheck(context, builder, ind, shapes[ax], ax)
|
|
output_indices.append(ind)
|
|
elif is_nonelike(idxty):
|
|
output_shapes.append(one)
|
|
output_strides.append(zero)
|
|
ax -= 1
|
|
else:
|
|
raise NotImplementedError("unexpected index type: %s" % (idxty,))
|
|
ax += 1
|
|
|
|
# Fill up missing dimensions at the end
|
|
assert ax <= aryty.ndim
|
|
while ax < aryty.ndim:
|
|
output_shapes.append(shapes[ax])
|
|
output_strides.append(strides[ax])
|
|
ax += 1
|
|
|
|
# No need to check wraparound, as negative indices were already
|
|
# fixed in the loop above.
|
|
dataptr = cgutils.get_item_pointer(context, builder, aryty, ary,
|
|
output_indices,
|
|
wraparound=False, boundscheck=False)
|
|
return (dataptr, output_shapes, output_strides)
|
|
|
|
|
|
def make_view(context, builder, aryty, ary, return_type,
|
|
data, shapes, strides):
|
|
"""
|
|
Build a view over the given array with the given parameters.
|
|
"""
|
|
retary = make_array(return_type)(context, builder)
|
|
populate_array(retary,
|
|
data=data,
|
|
shape=shapes,
|
|
strides=strides,
|
|
itemsize=ary.itemsize,
|
|
meminfo=ary.meminfo,
|
|
parent=ary.parent)
|
|
return retary
|
|
|
|
|
|
def _getitem_array_generic(context, builder, return_type, aryty, ary,
|
|
index_types, indices):
|
|
"""
|
|
Return the result of indexing *ary* with the given *indices*,
|
|
returning either a scalar or a view.
|
|
"""
|
|
dataptr, view_shapes, view_strides = \
|
|
basic_indexing(context, builder, aryty, ary, index_types, indices,
|
|
boundscheck=context.enable_boundscheck)
|
|
|
|
if isinstance(return_type, types.Buffer):
|
|
# Build array view
|
|
retary = make_view(context, builder, aryty, ary, return_type,
|
|
dataptr, view_shapes, view_strides)
|
|
return retary._getvalue()
|
|
else:
|
|
# Load scalar from 0-d result
|
|
assert not view_shapes
|
|
return load_item(context, builder, aryty, dataptr)
|
|
|
|
|
|
@lower_builtin(operator.getitem, types.Buffer, types.Integer)
|
|
@lower_builtin(operator.getitem, types.Buffer, types.SliceType)
|
|
def getitem_arraynd_intp(context, builder, sig, args):
|
|
"""
|
|
Basic indexing with an integer or a slice.
|
|
"""
|
|
aryty, idxty = sig.args
|
|
ary, idx = args
|
|
|
|
assert aryty.ndim >= 1
|
|
ary = make_array(aryty)(context, builder, ary)
|
|
|
|
res = _getitem_array_generic(context, builder, sig.return_type,
|
|
aryty, ary, (idxty,), (idx,))
|
|
return impl_ret_borrowed(context, builder, sig.return_type, res)
|
|
|
|
|
|
@lower_builtin(operator.getitem, types.Buffer, types.BaseTuple)
|
|
def getitem_array_tuple(context, builder, sig, args):
|
|
"""
|
|
Basic or advanced indexing with a tuple.
|
|
"""
|
|
aryty, tupty = sig.args
|
|
ary, tup = args
|
|
ary = make_array(aryty)(context, builder, ary)
|
|
|
|
index_types = tupty.types
|
|
indices = cgutils.unpack_tuple(builder, tup, count=len(tupty))
|
|
|
|
index_types, indices = normalize_indices(context, builder,
|
|
index_types, indices)
|
|
|
|
if any(isinstance(ty, types.Array) for ty in index_types):
|
|
# Advanced indexing
|
|
return fancy_getitem(context, builder, sig, args,
|
|
aryty, ary, index_types, indices)
|
|
|
|
res = _getitem_array_generic(context, builder, sig.return_type,
|
|
aryty, ary, index_types, indices)
|
|
return impl_ret_borrowed(context, builder, sig.return_type, res)
|
|
|
|
|
|
@lower_builtin(operator.setitem, types.Buffer, types.Any, types.Any)
|
|
def setitem_array(context, builder, sig, args):
|
|
"""
|
|
array[a] = scalar_or_array
|
|
array[a,..,b] = scalar_or_array
|
|
"""
|
|
aryty, idxty, valty = sig.args
|
|
ary, idx, val = args
|
|
|
|
if isinstance(idxty, types.BaseTuple):
|
|
index_types = idxty.types
|
|
indices = cgutils.unpack_tuple(builder, idx, count=len(idxty))
|
|
else:
|
|
index_types = (idxty,)
|
|
indices = (idx,)
|
|
|
|
ary = make_array(aryty)(context, builder, ary)
|
|
|
|
# First try basic indexing to see if a single array location is denoted.
|
|
index_types, indices = normalize_indices(context, builder,
|
|
index_types, indices)
|
|
try:
|
|
dataptr, shapes, strides = \
|
|
basic_indexing(context, builder, aryty, ary, index_types, indices,
|
|
boundscheck=context.enable_boundscheck)
|
|
except NotImplementedError:
|
|
use_fancy_indexing = True
|
|
else:
|
|
use_fancy_indexing = bool(shapes)
|
|
|
|
if use_fancy_indexing:
|
|
# Index describes a non-trivial view => use generic slice assignment
|
|
# (NOTE: this also handles scalar broadcasting)
|
|
return fancy_setslice(context, builder, sig, args,
|
|
index_types, indices)
|
|
|
|
# Store source value the given location
|
|
val = context.cast(builder, val, valty, aryty.dtype)
|
|
store_item(context, builder, aryty, val, dataptr)
|
|
|
|
|
|
@lower_builtin(len, types.Buffer)
|
|
def array_len(context, builder, sig, args):
|
|
(aryty,) = sig.args
|
|
(ary,) = args
|
|
arystty = make_array(aryty)
|
|
ary = arystty(context, builder, ary)
|
|
shapeary = ary.shape
|
|
res = builder.extract_value(shapeary, 0)
|
|
return impl_ret_untracked(context, builder, sig.return_type, res)
|
|
|
|
|
|
@lower_builtin("array.item", types.Array)
|
|
def array_item(context, builder, sig, args):
|
|
aryty, = sig.args
|
|
ary, = args
|
|
ary = make_array(aryty)(context, builder, ary)
|
|
|
|
nitems = ary.nitems
|
|
with builder.if_then(builder.icmp_signed('!=', nitems, nitems.type(1)),
|
|
likely=False):
|
|
msg = "item(): can only convert an array of size 1 to a Python scalar"
|
|
context.call_conv.return_user_exc(builder, ValueError, (msg,))
|
|
|
|
return load_item(context, builder, aryty, ary.data)
|
|
|
|
|
|
@lower_builtin("array.itemset", types.Array, types.Any)
|
|
def array_itemset(context, builder, sig, args):
|
|
aryty, valty = sig.args
|
|
ary, val = args
|
|
assert valty == aryty.dtype
|
|
ary = make_array(aryty)(context, builder, ary)
|
|
|
|
nitems = ary.nitems
|
|
with builder.if_then(builder.icmp_signed('!=', nitems, nitems.type(1)),
|
|
likely=False):
|
|
msg = "itemset(): can only write to an array of size 1"
|
|
context.call_conv.return_user_exc(builder, ValueError, (msg,))
|
|
|
|
store_item(context, builder, aryty, val, ary.data)
|
|
return context.get_dummy_value()
|
|
|
|
|
|
# ------------------------------------------------------------------------------
|
|
# Advanced / fancy indexing
|
|
|
|
|
|
class Indexer(object):
|
|
"""
|
|
Generic indexer interface, for generating indices over a fancy indexed
|
|
array on a single dimension.
|
|
"""
|
|
|
|
def prepare(self):
|
|
"""
|
|
Prepare the indexer by initializing any required variables, basic
|
|
blocks...
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def get_size(self):
|
|
"""
|
|
Return this dimension's size as an integer.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def get_shape(self):
|
|
"""
|
|
Return this dimension's shape as a tuple.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def get_index_bounds(self):
|
|
"""
|
|
Return a half-open [lower, upper) range of indices this dimension
|
|
is guaranteed not to step out of.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def loop_head(self):
|
|
"""
|
|
Start indexation loop. Return a (index, count) tuple.
|
|
*index* is an integer LLVM value representing the index over this
|
|
dimension.
|
|
*count* is either an integer LLVM value representing the current
|
|
iteration count, or None if this dimension should be omitted from
|
|
the indexation result.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def loop_tail(self):
|
|
"""
|
|
Finish indexation loop.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
|
|
class EntireIndexer(Indexer):
|
|
"""
|
|
Compute indices along an entire array dimension.
|
|
"""
|
|
|
|
def __init__(self, context, builder, aryty, ary, dim):
|
|
self.context = context
|
|
self.builder = builder
|
|
self.aryty = aryty
|
|
self.ary = ary
|
|
self.dim = dim
|
|
self.ll_intp = self.context.get_value_type(types.intp)
|
|
|
|
def prepare(self):
|
|
builder = self.builder
|
|
self.size = builder.extract_value(self.ary.shape, self.dim)
|
|
self.index = cgutils.alloca_once(builder, self.ll_intp)
|
|
self.bb_start = builder.append_basic_block()
|
|
self.bb_end = builder.append_basic_block()
|
|
|
|
def get_size(self):
|
|
return self.size
|
|
|
|
def get_shape(self):
|
|
return (self.size,)
|
|
|
|
def get_index_bounds(self):
|
|
# [0, size)
|
|
return (self.ll_intp(0), self.size)
|
|
|
|
def loop_head(self):
|
|
builder = self.builder
|
|
# Initialize loop variable
|
|
self.builder.store(Constant(self.ll_intp, 0), self.index)
|
|
builder.branch(self.bb_start)
|
|
builder.position_at_end(self.bb_start)
|
|
cur_index = builder.load(self.index)
|
|
with builder.if_then(builder.icmp_signed('>=', cur_index, self.size),
|
|
likely=False):
|
|
builder.branch(self.bb_end)
|
|
return cur_index, cur_index
|
|
|
|
def loop_tail(self):
|
|
builder = self.builder
|
|
next_index = cgutils.increment_index(builder, builder.load(self.index))
|
|
builder.store(next_index, self.index)
|
|
builder.branch(self.bb_start)
|
|
builder.position_at_end(self.bb_end)
|
|
|
|
|
|
class IntegerIndexer(Indexer):
|
|
"""
|
|
Compute indices from a single integer.
|
|
"""
|
|
|
|
def __init__(self, context, builder, idx):
|
|
self.context = context
|
|
self.builder = builder
|
|
self.idx = idx
|
|
self.ll_intp = self.context.get_value_type(types.intp)
|
|
|
|
def prepare(self):
|
|
pass
|
|
|
|
def get_size(self):
|
|
return Constant(self.ll_intp, 1)
|
|
|
|
def get_shape(self):
|
|
return ()
|
|
|
|
def get_index_bounds(self):
|
|
# [idx, idx+1)
|
|
return (self.idx, self.builder.add(self.idx, self.get_size()))
|
|
|
|
def loop_head(self):
|
|
return self.idx, None
|
|
|
|
def loop_tail(self):
|
|
pass
|
|
|
|
|
|
class IntegerArrayIndexer(Indexer):
|
|
"""
|
|
Compute indices from an array of integer indices.
|
|
"""
|
|
|
|
def __init__(self, context, builder, idxty, idxary, size):
|
|
self.context = context
|
|
self.builder = builder
|
|
self.idxty = idxty
|
|
self.idxary = idxary
|
|
self.size = size
|
|
assert idxty.ndim == 1
|
|
self.ll_intp = self.context.get_value_type(types.intp)
|
|
|
|
def prepare(self):
|
|
builder = self.builder
|
|
self.idx_size = cgutils.unpack_tuple(builder, self.idxary.shape)[0]
|
|
self.idx_index = cgutils.alloca_once(builder, self.ll_intp)
|
|
self.bb_start = builder.append_basic_block()
|
|
self.bb_end = builder.append_basic_block()
|
|
|
|
def get_size(self):
|
|
return self.idx_size
|
|
|
|
def get_shape(self):
|
|
return (self.idx_size,)
|
|
|
|
def get_index_bounds(self):
|
|
# Pessimal heuristic, as we don't want to scan for the min and max
|
|
return (self.ll_intp(0), self.size)
|
|
|
|
def loop_head(self):
|
|
builder = self.builder
|
|
# Initialize loop variable
|
|
self.builder.store(Constant(self.ll_intp, 0), self.idx_index)
|
|
builder.branch(self.bb_start)
|
|
builder.position_at_end(self.bb_start)
|
|
cur_index = builder.load(self.idx_index)
|
|
with builder.if_then(
|
|
builder.icmp_signed('>=', cur_index, self.idx_size),
|
|
likely=False
|
|
):
|
|
builder.branch(self.bb_end)
|
|
# Load the actual index from the array of indices
|
|
index = _getitem_array_single_int(
|
|
self.context, builder, self.idxty.dtype, self.idxty, self.idxary,
|
|
cur_index
|
|
)
|
|
index = fix_integer_index(self.context, builder,
|
|
self.idxty.dtype, index, self.size)
|
|
return index, cur_index
|
|
|
|
def loop_tail(self):
|
|
builder = self.builder
|
|
next_index = cgutils.increment_index(builder,
|
|
builder.load(self.idx_index))
|
|
builder.store(next_index, self.idx_index)
|
|
builder.branch(self.bb_start)
|
|
builder.position_at_end(self.bb_end)
|
|
|
|
|
|
class BooleanArrayIndexer(Indexer):
|
|
"""
|
|
Compute indices from an array of boolean predicates.
|
|
"""
|
|
|
|
def __init__(self, context, builder, idxty, idxary):
|
|
self.context = context
|
|
self.builder = builder
|
|
self.idxty = idxty
|
|
self.idxary = idxary
|
|
assert idxty.ndim == 1
|
|
self.ll_intp = self.context.get_value_type(types.intp)
|
|
self.zero = Constant(self.ll_intp, 0)
|
|
|
|
def prepare(self):
|
|
builder = self.builder
|
|
self.size = cgutils.unpack_tuple(builder, self.idxary.shape)[0]
|
|
self.idx_index = cgutils.alloca_once(builder, self.ll_intp)
|
|
self.count = cgutils.alloca_once(builder, self.ll_intp)
|
|
self.bb_start = builder.append_basic_block()
|
|
self.bb_tail = builder.append_basic_block()
|
|
self.bb_end = builder.append_basic_block()
|
|
|
|
def get_size(self):
|
|
builder = self.builder
|
|
count = cgutils.alloca_once_value(builder, self.zero)
|
|
# Sum all true values
|
|
with cgutils.for_range(builder, self.size) as loop:
|
|
c = builder.load(count)
|
|
pred = _getitem_array_single_int(
|
|
self.context, builder, self.idxty.dtype,
|
|
self.idxty, self.idxary, loop.index
|
|
)
|
|
c = builder.add(c, builder.zext(pred, c.type))
|
|
builder.store(c, count)
|
|
|
|
return builder.load(count)
|
|
|
|
def get_shape(self):
|
|
return (self.get_size(),)
|
|
|
|
def get_index_bounds(self):
|
|
# Pessimal heuristic, as we don't want to scan for the
|
|
# first and last true items
|
|
return (self.ll_intp(0), self.size)
|
|
|
|
def loop_head(self):
|
|
builder = self.builder
|
|
# Initialize loop variable
|
|
self.builder.store(self.zero, self.idx_index)
|
|
self.builder.store(self.zero, self.count)
|
|
builder.branch(self.bb_start)
|
|
builder.position_at_end(self.bb_start)
|
|
cur_index = builder.load(self.idx_index)
|
|
cur_count = builder.load(self.count)
|
|
with builder.if_then(builder.icmp_signed('>=', cur_index, self.size),
|
|
likely=False):
|
|
builder.branch(self.bb_end)
|
|
# Load the predicate and branch if false
|
|
pred = _getitem_array_single_int(
|
|
self.context, builder, self.idxty.dtype, self.idxty, self.idxary,
|
|
cur_index
|
|
)
|
|
with builder.if_then(builder.not_(pred)):
|
|
builder.branch(self.bb_tail)
|
|
# Increment the count for next iteration
|
|
next_count = cgutils.increment_index(builder, cur_count)
|
|
builder.store(next_count, self.count)
|
|
return cur_index, cur_count
|
|
|
|
def loop_tail(self):
|
|
builder = self.builder
|
|
builder.branch(self.bb_tail)
|
|
builder.position_at_end(self.bb_tail)
|
|
next_index = cgutils.increment_index(builder,
|
|
builder.load(self.idx_index))
|
|
builder.store(next_index, self.idx_index)
|
|
builder.branch(self.bb_start)
|
|
builder.position_at_end(self.bb_end)
|
|
|
|
|
|
class SliceIndexer(Indexer):
|
|
"""
|
|
Compute indices along a slice.
|
|
"""
|
|
|
|
def __init__(self, context, builder, aryty, ary, dim, idxty, slice):
|
|
self.context = context
|
|
self.builder = builder
|
|
self.aryty = aryty
|
|
self.ary = ary
|
|
self.dim = dim
|
|
self.idxty = idxty
|
|
self.slice = slice
|
|
self.ll_intp = self.context.get_value_type(types.intp)
|
|
self.zero = Constant(self.ll_intp, 0)
|
|
|
|
def prepare(self):
|
|
builder = self.builder
|
|
# Fix slice for the dimension's size
|
|
self.dim_size = builder.extract_value(self.ary.shape, self.dim)
|
|
slicing.guard_invalid_slice(self.context, builder, self.idxty,
|
|
self.slice)
|
|
slicing.fix_slice(builder, self.slice, self.dim_size)
|
|
self.is_step_negative = cgutils.is_neg_int(builder, self.slice.step)
|
|
# Create loop entities
|
|
self.index = cgutils.alloca_once(builder, self.ll_intp)
|
|
self.count = cgutils.alloca_once(builder, self.ll_intp)
|
|
self.bb_start = builder.append_basic_block()
|
|
self.bb_end = builder.append_basic_block()
|
|
|
|
def get_size(self):
|
|
return slicing.get_slice_length(self.builder, self.slice)
|
|
|
|
def get_shape(self):
|
|
return (self.get_size(),)
|
|
|
|
def get_index_bounds(self):
|
|
lower, upper = slicing.get_slice_bounds(self.builder, self.slice)
|
|
return lower, upper
|
|
|
|
def loop_head(self):
|
|
builder = self.builder
|
|
# Initialize loop variable
|
|
self.builder.store(self.slice.start, self.index)
|
|
self.builder.store(self.zero, self.count)
|
|
builder.branch(self.bb_start)
|
|
builder.position_at_end(self.bb_start)
|
|
cur_index = builder.load(self.index)
|
|
cur_count = builder.load(self.count)
|
|
is_finished = builder.select(self.is_step_negative,
|
|
builder.icmp_signed('<=', cur_index,
|
|
self.slice.stop),
|
|
builder.icmp_signed('>=', cur_index,
|
|
self.slice.stop))
|
|
with builder.if_then(is_finished, likely=False):
|
|
builder.branch(self.bb_end)
|
|
return cur_index, cur_count
|
|
|
|
def loop_tail(self):
|
|
builder = self.builder
|
|
next_index = builder.add(builder.load(self.index), self.slice.step,
|
|
flags=['nsw'])
|
|
builder.store(next_index, self.index)
|
|
next_count = cgutils.increment_index(builder, builder.load(self.count))
|
|
builder.store(next_count, self.count)
|
|
builder.branch(self.bb_start)
|
|
builder.position_at_end(self.bb_end)
|
|
|
|
|
|
class FancyIndexer(object):
|
|
"""
|
|
Perform fancy indexing on the given array.
|
|
"""
|
|
|
|
def __init__(self, context, builder, aryty, ary, index_types, indices):
|
|
self.context = context
|
|
self.builder = builder
|
|
self.aryty = aryty
|
|
self.shapes = cgutils.unpack_tuple(builder, ary.shape, aryty.ndim)
|
|
self.strides = cgutils.unpack_tuple(builder, ary.strides, aryty.ndim)
|
|
self.ll_intp = self.context.get_value_type(types.intp)
|
|
self.newaxes = []
|
|
|
|
indexers = []
|
|
num_newaxes = len([idx for idx in index_types if is_nonelike(idx)])
|
|
|
|
ax = 0 # keeps track of position of original axes
|
|
new_ax = 0 # keeps track of position for inserting new axes
|
|
for indexval, idxty in zip(indices, index_types):
|
|
if idxty is types.ellipsis:
|
|
# Fill up missing dimensions at the middle
|
|
n_missing = aryty.ndim - len(indices) + 1 + num_newaxes
|
|
for i in range(n_missing):
|
|
indexer = EntireIndexer(context, builder, aryty, ary, ax)
|
|
indexers.append(indexer)
|
|
ax += 1
|
|
new_ax += 1
|
|
continue
|
|
|
|
# Regular index value
|
|
if isinstance(idxty, types.SliceType):
|
|
slice = context.make_helper(builder, idxty, indexval)
|
|
indexer = SliceIndexer(context, builder, aryty, ary, ax,
|
|
idxty, slice)
|
|
indexers.append(indexer)
|
|
elif isinstance(idxty, types.Integer):
|
|
ind = fix_integer_index(context, builder, idxty, indexval,
|
|
self.shapes[ax])
|
|
indexer = IntegerIndexer(context, builder, ind)
|
|
indexers.append(indexer)
|
|
elif isinstance(idxty, types.Array):
|
|
idxary = make_array(idxty)(context, builder, indexval)
|
|
if isinstance(idxty.dtype, types.Integer):
|
|
indexer = IntegerArrayIndexer(context, builder,
|
|
idxty, idxary,
|
|
self.shapes[ax])
|
|
elif isinstance(idxty.dtype, types.Boolean):
|
|
indexer = BooleanArrayIndexer(context, builder,
|
|
idxty, idxary)
|
|
else:
|
|
assert 0
|
|
indexers.append(indexer)
|
|
elif is_nonelike(idxty):
|
|
self.newaxes.append(new_ax)
|
|
ax -= 1
|
|
else:
|
|
raise AssertionError("unexpected index type: %s" % (idxty,))
|
|
ax += 1
|
|
new_ax += 1
|
|
|
|
# Fill up missing dimensions at the end
|
|
assert ax <= aryty.ndim, (ax, aryty.ndim)
|
|
while ax < aryty.ndim:
|
|
indexer = EntireIndexer(context, builder, aryty, ary, ax)
|
|
indexers.append(indexer)
|
|
ax += 1
|
|
|
|
assert len(indexers) == aryty.ndim, (len(indexers), aryty.ndim)
|
|
self.indexers = indexers
|
|
|
|
def prepare(self):
|
|
for i in self.indexers:
|
|
i.prepare()
|
|
|
|
one = self.context.get_constant(types.intp, 1)
|
|
|
|
# Compute the resulting shape given by the indices
|
|
res_shape = [i.get_shape() for i in self.indexers]
|
|
|
|
# At every position where newaxis/None is present insert
|
|
# one as a constant shape in the resulting list of shapes.
|
|
for i in self.newaxes:
|
|
res_shape.insert(i, (one,))
|
|
|
|
# Store the shape as a tuple, we can't do a simple
|
|
# tuple(res_shape) here since res_shape is a list
|
|
# of tuples which may be differently sized.
|
|
self.indexers_shape = sum(res_shape, ())
|
|
|
|
def get_shape(self):
|
|
"""
|
|
Get the resulting data shape as Python tuple.
|
|
"""
|
|
return self.indexers_shape
|
|
|
|
def get_offset_bounds(self, strides, itemsize):
|
|
"""
|
|
Get a half-open [lower, upper) range of byte offsets spanned by
|
|
the indexer with the given strides and itemsize. The indexer is
|
|
guaranteed to not go past those bounds.
|
|
"""
|
|
assert len(strides) == self.aryty.ndim
|
|
builder = self.builder
|
|
is_empty = cgutils.false_bit
|
|
zero = self.ll_intp(0)
|
|
one = self.ll_intp(1)
|
|
lower = zero
|
|
upper = zero
|
|
for indexer, shape, stride in zip(self.indexers, self.indexers_shape,
|
|
strides):
|
|
is_empty = builder.or_(is_empty,
|
|
builder.icmp_unsigned('==', shape, zero))
|
|
# Compute [lower, upper) indices on this dimension
|
|
lower_index, upper_index = indexer.get_index_bounds()
|
|
lower_offset = builder.mul(stride, lower_index)
|
|
upper_offset = builder.mul(stride, builder.sub(upper_index, one))
|
|
# Adjust total interval
|
|
is_downwards = builder.icmp_signed('<', stride, zero)
|
|
lower = builder.add(lower,
|
|
builder.select(is_downwards,
|
|
upper_offset,
|
|
lower_offset))
|
|
upper = builder.add(upper,
|
|
builder.select(is_downwards,
|
|
lower_offset,
|
|
upper_offset))
|
|
# Make interval half-open
|
|
upper = builder.add(upper, itemsize)
|
|
# Adjust for empty shape
|
|
lower = builder.select(is_empty, zero, lower)
|
|
upper = builder.select(is_empty, zero, upper)
|
|
return lower, upper
|
|
|
|
def begin_loops(self):
|
|
indices, counts = zip(*(i.loop_head() for i in self.indexers))
|
|
return indices, counts
|
|
|
|
def end_loops(self):
|
|
for i in reversed(self.indexers):
|
|
i.loop_tail()
|
|
|
|
|
|
def fancy_getitem(context, builder, sig, args,
|
|
aryty, ary, index_types, indices):
|
|
|
|
shapes = cgutils.unpack_tuple(builder, ary.shape)
|
|
strides = cgutils.unpack_tuple(builder, ary.strides)
|
|
data = ary.data
|
|
|
|
indexer = FancyIndexer(context, builder, aryty, ary,
|
|
index_types, indices)
|
|
indexer.prepare()
|
|
|
|
# Construct output array
|
|
out_ty = sig.return_type
|
|
out_shapes = indexer.get_shape()
|
|
|
|
out = _empty_nd_impl(context, builder, out_ty, out_shapes)
|
|
out_data = out.data
|
|
out_idx = cgutils.alloca_once_value(builder,
|
|
context.get_constant(types.intp, 0))
|
|
|
|
# Loop on source and copy to destination
|
|
indices, _ = indexer.begin_loops()
|
|
|
|
# No need to check for wraparound, as the indexers all ensure
|
|
# a positive index is returned.
|
|
ptr = cgutils.get_item_pointer2(context, builder, data, shapes, strides,
|
|
aryty.layout, indices, wraparound=False,
|
|
boundscheck=context.enable_boundscheck)
|
|
val = load_item(context, builder, aryty, ptr)
|
|
|
|
# Since the destination is C-contiguous, no need for multi-dimensional
|
|
# indexing.
|
|
cur = builder.load(out_idx)
|
|
ptr = builder.gep(out_data, [cur])
|
|
store_item(context, builder, out_ty, val, ptr)
|
|
next_idx = cgutils.increment_index(builder, cur)
|
|
builder.store(next_idx, out_idx)
|
|
|
|
indexer.end_loops()
|
|
|
|
return impl_ret_new_ref(context, builder, out_ty, out._getvalue())
|
|
|
|
|
|
@lower_builtin(operator.getitem, types.Buffer, types.Array)
|
|
def fancy_getitem_array(context, builder, sig, args):
|
|
"""
|
|
Advanced or basic indexing with an array.
|
|
"""
|
|
aryty, idxty = sig.args
|
|
ary, idx = args
|
|
ary = make_array(aryty)(context, builder, ary)
|
|
if idxty.ndim == 0:
|
|
# 0-d array index acts as a basic integer index
|
|
idxty, idx = normalize_index(context, builder, idxty, idx)
|
|
res = _getitem_array_generic(context, builder, sig.return_type,
|
|
aryty, ary, (idxty,), (idx,))
|
|
return impl_ret_borrowed(context, builder, sig.return_type, res)
|
|
else:
|
|
# Advanced indexing
|
|
return fancy_getitem(context, builder, sig, args,
|
|
aryty, ary, (idxty,), (idx,))
|
|
|
|
|
|
def offset_bounds_from_strides(context, builder, arrty, arr, shapes, strides):
|
|
"""
|
|
Compute a half-open range [lower, upper) of byte offsets from the
|
|
array's data pointer, that bound the in-memory extent of the array.
|
|
|
|
This mimics offset_bounds_from_strides() from
|
|
numpy/core/src/private/mem_overlap.c
|
|
"""
|
|
itemsize = arr.itemsize
|
|
zero = itemsize.type(0)
|
|
one = zero.type(1)
|
|
if arrty.layout in 'CF':
|
|
# Array is contiguous: contents are laid out sequentially
|
|
# starting from arr.data and upwards
|
|
lower = zero
|
|
upper = builder.mul(itemsize, arr.nitems)
|
|
else:
|
|
# Non-contiguous array: need to examine strides
|
|
lower = zero
|
|
upper = zero
|
|
for i in range(arrty.ndim):
|
|
# Compute the largest byte offset on this dimension
|
|
# max_axis_offset = strides[i] * (shapes[i] - 1)
|
|
# (shapes[i] == 0 is catered for by the empty array case below)
|
|
max_axis_offset = builder.mul(strides[i],
|
|
builder.sub(shapes[i], one))
|
|
is_upwards = builder.icmp_signed('>=', max_axis_offset, zero)
|
|
# Expand either upwards or downwards depending on stride
|
|
upper = builder.select(is_upwards,
|
|
builder.add(upper, max_axis_offset), upper)
|
|
lower = builder.select(is_upwards,
|
|
lower, builder.add(lower, max_axis_offset))
|
|
# Return a half-open range
|
|
upper = builder.add(upper, itemsize)
|
|
# Adjust for empty arrays
|
|
is_empty = builder.icmp_signed('==', arr.nitems, zero)
|
|
upper = builder.select(is_empty, zero, upper)
|
|
lower = builder.select(is_empty, zero, lower)
|
|
|
|
return lower, upper
|
|
|
|
|
|
def compute_memory_extents(context, builder, lower, upper, data):
|
|
"""
|
|
Given [lower, upper) byte offsets and a base data pointer,
|
|
compute the memory pointer bounds as pointer-sized integers.
|
|
"""
|
|
data_ptr_as_int = builder.ptrtoint(data, lower.type)
|
|
start = builder.add(data_ptr_as_int, lower)
|
|
end = builder.add(data_ptr_as_int, upper)
|
|
return start, end
|
|
|
|
|
|
def get_array_memory_extents(context, builder, arrty, arr, shapes, strides,
|
|
data):
|
|
"""
|
|
Compute a half-open range [start, end) of pointer-sized integers
|
|
which fully contain the array data.
|
|
"""
|
|
lower, upper = offset_bounds_from_strides(context, builder, arrty, arr,
|
|
shapes, strides)
|
|
return compute_memory_extents(context, builder, lower, upper, data)
|
|
|
|
|
|
def extents_may_overlap(context, builder, a_start, a_end, b_start, b_end):
|
|
"""
|
|
Whether two memory extents [a_start, a_end) and [b_start, b_end)
|
|
may overlap.
|
|
"""
|
|
# Comparisons are unsigned, since we are really comparing pointers
|
|
may_overlap = builder.and_(
|
|
builder.icmp_unsigned('<', a_start, b_end),
|
|
builder.icmp_unsigned('<', b_start, a_end),
|
|
)
|
|
return may_overlap
|
|
|
|
|
|
def maybe_copy_source(context, builder, use_copy,
|
|
srcty, src, src_shapes, src_strides, src_data):
|
|
ptrty = src_data.type
|
|
|
|
copy_layout = 'C'
|
|
copy_data = cgutils.alloca_once_value(builder, src_data)
|
|
copy_shapes = src_shapes
|
|
copy_strides = None # unneeded for contiguous arrays
|
|
|
|
with builder.if_then(use_copy, likely=False):
|
|
# Allocate temporary scratchpad
|
|
# XXX: should we use a stack-allocated array for very small
|
|
# data sizes?
|
|
allocsize = builder.mul(src.itemsize, src.nitems)
|
|
data = context.nrt.allocate(builder, allocsize)
|
|
voidptrty = data.type
|
|
data = builder.bitcast(data, ptrty)
|
|
builder.store(data, copy_data)
|
|
|
|
# Copy source data into scratchpad
|
|
intp_t = context.get_value_type(types.intp)
|
|
|
|
with cgutils.loop_nest(builder, src_shapes, intp_t) as indices:
|
|
src_ptr = cgutils.get_item_pointer2(context, builder, src_data,
|
|
src_shapes, src_strides,
|
|
srcty.layout, indices)
|
|
dest_ptr = cgutils.get_item_pointer2(context, builder, data,
|
|
copy_shapes, copy_strides,
|
|
copy_layout, indices)
|
|
builder.store(builder.load(src_ptr), dest_ptr)
|
|
|
|
def src_getitem(source_indices):
|
|
src_ptr = cgutils.alloca_once(builder, ptrty)
|
|
with builder.if_else(use_copy, likely=False) as (if_copy, otherwise):
|
|
with if_copy:
|
|
builder.store(
|
|
cgutils.get_item_pointer2(context, builder,
|
|
builder.load(copy_data),
|
|
copy_shapes, copy_strides,
|
|
copy_layout, source_indices,
|
|
wraparound=False),
|
|
src_ptr)
|
|
with otherwise:
|
|
builder.store(
|
|
cgutils.get_item_pointer2(context, builder, src_data,
|
|
src_shapes, src_strides,
|
|
srcty.layout, source_indices,
|
|
wraparound=False),
|
|
src_ptr)
|
|
return load_item(context, builder, srcty, builder.load(src_ptr))
|
|
|
|
def src_cleanup():
|
|
# Deallocate memory
|
|
with builder.if_then(use_copy, likely=False):
|
|
data = builder.load(copy_data)
|
|
data = builder.bitcast(data, voidptrty)
|
|
context.nrt.free(builder, data)
|
|
|
|
return src_getitem, src_cleanup
|
|
|
|
|
|
def _bc_adjust_dimension(context, builder, shapes, strides, target_shape):
|
|
"""
|
|
Preprocess dimension for broadcasting.
|
|
Returns (shapes, strides) such that the ndim match *target_shape*.
|
|
When expanding to higher ndim, the returning shapes and strides are
|
|
prepended with ones and zeros, respectively.
|
|
When truncating to lower ndim, the shapes are checked (in runtime).
|
|
All extra dimension must have size of 1.
|
|
"""
|
|
zero = context.get_constant(types.uintp, 0)
|
|
one = context.get_constant(types.uintp, 1)
|
|
|
|
# Adjust for broadcasting to higher dimension
|
|
if len(target_shape) > len(shapes):
|
|
nd_diff = len(target_shape) - len(shapes)
|
|
# Fill missing shapes with one, strides with zeros
|
|
shapes = [one] * nd_diff + shapes
|
|
strides = [zero] * nd_diff + strides
|
|
# Adjust for broadcasting to lower dimension
|
|
elif len(target_shape) < len(shapes):
|
|
# Accepted if all extra dims has shape 1
|
|
nd_diff = len(shapes) - len(target_shape)
|
|
dim_is_one = [builder.icmp_unsigned('==', sh, one)
|
|
for sh in shapes[:nd_diff]]
|
|
accepted = functools.reduce(builder.and_, dim_is_one,
|
|
cgutils.true_bit)
|
|
# Check error
|
|
with builder.if_then(builder.not_(accepted), likely=False):
|
|
msg = "cannot broadcast source array for assignment"
|
|
context.call_conv.return_user_exc(builder, ValueError, (msg,))
|
|
# Truncate extra shapes, strides
|
|
shapes = shapes[nd_diff:]
|
|
strides = strides[nd_diff:]
|
|
|
|
return shapes, strides
|
|
|
|
|
|
def _bc_adjust_shape_strides(context, builder, shapes, strides, target_shape):
|
|
"""
|
|
Broadcast shapes and strides to target_shape given that their ndim already
|
|
matches. For each location where the shape is 1 and does not match the
|
|
dim for target, it is set to the value at the target and the stride is
|
|
set to zero.
|
|
"""
|
|
bc_shapes = []
|
|
bc_strides = []
|
|
zero = context.get_constant(types.uintp, 0)
|
|
one = context.get_constant(types.uintp, 1)
|
|
# Adjust all mismatching ones in shape
|
|
mismatch = [builder.icmp_signed('!=', tar, old)
|
|
for tar, old in zip(target_shape, shapes)]
|
|
src_is_one = [builder.icmp_signed('==', old, one) for old in shapes]
|
|
preds = [builder.and_(x, y) for x, y in zip(mismatch, src_is_one)]
|
|
bc_shapes = [builder.select(p, tar, old)
|
|
for p, tar, old in zip(preds, target_shape, shapes)]
|
|
bc_strides = [builder.select(p, zero, old)
|
|
for p, old in zip(preds, strides)]
|
|
return bc_shapes, bc_strides
|
|
|
|
|
|
def _broadcast_to_shape(context, builder, arrtype, arr, target_shape):
|
|
"""
|
|
Broadcast the given array to the target_shape.
|
|
Returns (array_type, array)
|
|
"""
|
|
# Compute broadcasted shape and strides
|
|
shapes = cgutils.unpack_tuple(builder, arr.shape)
|
|
strides = cgutils.unpack_tuple(builder, arr.strides)
|
|
|
|
shapes, strides = _bc_adjust_dimension(context, builder, shapes, strides,
|
|
target_shape)
|
|
shapes, strides = _bc_adjust_shape_strides(context, builder, shapes,
|
|
strides, target_shape)
|
|
new_arrtype = arrtype.copy(ndim=len(target_shape), layout='A')
|
|
# Create new view
|
|
new_arr = make_array(new_arrtype)(context, builder)
|
|
populate_array(new_arr,
|
|
data=arr.data,
|
|
shape=cgutils.pack_array(builder, shapes),
|
|
strides=cgutils.pack_array(builder, strides),
|
|
itemsize=arr.itemsize,
|
|
meminfo=arr.meminfo,
|
|
parent=arr.parent)
|
|
return new_arrtype, new_arr
|
|
|
|
|
|
@intrinsic
|
|
def _numpy_broadcast_to(typingctx, array, shape):
|
|
ret = array.copy(ndim=shape.count, layout='A', readonly=True)
|
|
sig = ret(array, shape)
|
|
|
|
def codegen(context, builder, sig, args):
|
|
src, shape_ = args
|
|
srcty = sig.args[0]
|
|
|
|
src = make_array(srcty)(context, builder, src)
|
|
shape_ = cgutils.unpack_tuple(builder, shape_)
|
|
_, dest = _broadcast_to_shape(context, builder, srcty, src, shape_,)
|
|
|
|
# Hack to get np.broadcast_to to return a read-only array
|
|
setattr(dest, 'parent', Constant(
|
|
context.get_value_type(dest._datamodel.get_type('parent')),
|
|
None))
|
|
|
|
res = dest._getvalue()
|
|
return impl_ret_borrowed(context, builder, sig.return_type, res)
|
|
return sig, codegen
|
|
|
|
|
|
@intrinsic
|
|
def get_readonly_array(typingctx, arr):
|
|
# returns a copy of arr which is readonly
|
|
ret = arr.copy(readonly=True)
|
|
sig = ret(arr)
|
|
|
|
def codegen(context, builder, sig, args):
|
|
[src] = args
|
|
srcty = sig.args[0]
|
|
|
|
dest = make_array(srcty)(context, builder, src)
|
|
# Hack to return a read-only array
|
|
dest.parent = cgutils.get_null_value(dest.parent.type)
|
|
res = dest._getvalue()
|
|
return impl_ret_borrowed(context, builder, sig.return_type, res)
|
|
return sig, codegen
|
|
|
|
|
|
@register_jitable
|
|
def _can_broadcast(array, dest_shape):
|
|
src_shape = array.shape
|
|
src_ndim = len(src_shape)
|
|
dest_ndim = len(dest_shape)
|
|
if src_ndim > dest_ndim:
|
|
raise ValueError('input operand has more dimensions than allowed '
|
|
'by the axis remapping')
|
|
for size in dest_shape:
|
|
if size < 0:
|
|
raise ValueError('all elements of broadcast shape must be '
|
|
'non-negative')
|
|
|
|
# based on _broadcast_onto function in numba/np/npyimpl.py
|
|
src_index = 0
|
|
dest_index = dest_ndim - src_ndim
|
|
while src_index < src_ndim:
|
|
src_dim = src_shape[src_index]
|
|
dest_dim = dest_shape[dest_index]
|
|
# possible cases for (src_dim, dest_dim):
|
|
# * (1, 1) -> Ok
|
|
# * (>1, 1) -> Error!
|
|
# * (>1, >1) -> src_dim == dest_dim else error!
|
|
# * (1, >1) -> Ok
|
|
if src_dim == dest_dim or src_dim == 1:
|
|
src_index += 1
|
|
dest_index += 1
|
|
else:
|
|
raise ValueError('operands could not be broadcast together '
|
|
'with remapped shapes')
|
|
|
|
|
|
def _default_broadcast_to_impl(array, shape):
|
|
array = np.asarray(array)
|
|
_can_broadcast(array, shape)
|
|
return _numpy_broadcast_to(array, shape)
|
|
|
|
|
|
@overload(np.broadcast_to)
|
|
def numpy_broadcast_to(array, shape):
|
|
if not type_can_asarray(array):
|
|
raise errors.TypingError('The first argument "array" must '
|
|
'be array-like')
|
|
|
|
if isinstance(shape, types.Integer):
|
|
def impl(array, shape):
|
|
return np.broadcast_to(array, (shape,))
|
|
return impl
|
|
|
|
elif isinstance(shape, types.UniTuple):
|
|
if not isinstance(shape.dtype, types.Integer):
|
|
msg = 'The second argument "shape" must be a tuple of integers'
|
|
raise errors.TypingError(msg)
|
|
return _default_broadcast_to_impl
|
|
|
|
elif isinstance(shape, types.Tuple) and shape.count > 0:
|
|
# check if all types are integers
|
|
if not all([isinstance(typ, types.IntegerLiteral) for typ in shape]):
|
|
msg = f'"{shape}" object cannot be interpreted as an integer'
|
|
raise errors.TypingError(msg)
|
|
return _default_broadcast_to_impl
|
|
elif isinstance(shape, types.Tuple) and shape.count == 0:
|
|
is_scalar_array = isinstance(array, types.Array) and array.ndim == 0
|
|
if type_is_scalar(array) or is_scalar_array:
|
|
|
|
def impl(array, shape): # broadcast_to(array, ())
|
|
# Array type must be supported by "type_can_asarray"
|
|
# Quick note that unicode types are not supported!
|
|
array = np.asarray(array)
|
|
return get_readonly_array(array)
|
|
return impl
|
|
|
|
else:
|
|
msg = 'Cannot broadcast a non-scalar to a scalar array'
|
|
raise errors.TypingError(msg)
|
|
else:
|
|
msg = ('The argument "shape" must be a tuple or an integer. '
|
|
'Got %s' % shape)
|
|
raise errors.TypingError(msg)
|
|
|
|
|
|
@register_jitable
|
|
def numpy_broadcast_shapes_list(r, m, shape):
|
|
for i in range(len(shape)):
|
|
k = m - len(shape) + i
|
|
tmp = shape[i]
|
|
if tmp < 0:
|
|
raise ValueError("negative dimensions are not allowed")
|
|
if tmp == 1:
|
|
continue
|
|
if r[k] == 1:
|
|
r[k] = tmp
|
|
elif r[k] != tmp:
|
|
raise ValueError("shape mismatch: objects"
|
|
" cannot be broadcast"
|
|
" to a single shape")
|
|
|
|
|
|
@overload(np.broadcast_shapes)
|
|
def ol_numpy_broadcast_shapes(*args):
|
|
# Based on https://github.com/numpy/numpy/blob/f702b26fff3271ba6a6ba29a021fc19051d1f007/numpy/core/src/multiarray/iterators.c#L1129-L1212 # noqa
|
|
for idx, arg in enumerate(args):
|
|
is_int = isinstance(arg, types.Integer)
|
|
is_int_tuple = isinstance(arg, types.UniTuple) and \
|
|
isinstance(arg.dtype, types.Integer)
|
|
is_empty_tuple = isinstance(arg, types.Tuple) and len(arg.types) == 0
|
|
if not (is_int or is_int_tuple or is_empty_tuple):
|
|
msg = (f'Argument {idx} must be either an int or tuple[int]. '
|
|
f'Got {arg}')
|
|
raise errors.TypingError(msg)
|
|
|
|
# discover the number of dimensions
|
|
m = 0
|
|
for arg in args:
|
|
if isinstance(arg, types.Integer):
|
|
m = max(m, 1)
|
|
elif isinstance(arg, types.BaseTuple):
|
|
m = max(m, len(arg))
|
|
|
|
if m == 0:
|
|
return lambda *args: ()
|
|
else:
|
|
tup_init = (1,) * m
|
|
|
|
def impl(*args):
|
|
# propagate args
|
|
r = [1] * m
|
|
tup = tup_init
|
|
for arg in literal_unroll(args):
|
|
if isinstance(arg, tuple) and len(arg) > 0:
|
|
numpy_broadcast_shapes_list(r, m, arg)
|
|
elif isinstance(arg, int):
|
|
numpy_broadcast_shapes_list(r, m, (arg,))
|
|
for idx, elem in enumerate(r):
|
|
tup = tuple_setitem(tup, idx, elem)
|
|
return tup
|
|
return impl
|
|
|
|
|
|
@overload(np.broadcast_arrays)
|
|
def numpy_broadcast_arrays(*args):
|
|
|
|
for idx, arg in enumerate(args):
|
|
if not type_can_asarray(arg):
|
|
raise errors.TypingError(f'Argument "{idx}" must '
|
|
'be array-like')
|
|
|
|
unified_dtype = None
|
|
dt = None
|
|
for arg in args:
|
|
if isinstance(arg, (types.Array, types.BaseTuple)):
|
|
dt = arg.dtype
|
|
else:
|
|
dt = arg
|
|
|
|
if unified_dtype is None:
|
|
unified_dtype = dt
|
|
elif unified_dtype != dt:
|
|
raise errors.TypingError('Mismatch of argument types. Numba cannot '
|
|
'broadcast arrays with different types. '
|
|
f'Got {args}')
|
|
|
|
# number of dimensions
|
|
m = 0
|
|
for idx, arg in enumerate(args):
|
|
if isinstance(arg, types.ArrayCompatible):
|
|
m = max(m, arg.ndim)
|
|
elif isinstance(arg, (types.Number, types.Boolean, types.BaseTuple)):
|
|
m = max(m, 1)
|
|
else:
|
|
raise errors.TypingError(f'Unhandled type {arg}')
|
|
|
|
tup_init = (0,) * m
|
|
|
|
def impl(*args):
|
|
# find out the output shape
|
|
# we can't call np.broadcast_shapes here since args may have arrays
|
|
# with different shapes and it is not possible to create a list
|
|
# with those shapes dynamically
|
|
shape = [1] * m
|
|
for array in literal_unroll(args):
|
|
numpy_broadcast_shapes_list(shape, m, np.asarray(array).shape)
|
|
|
|
tup = tup_init
|
|
|
|
for i in range(m):
|
|
tup = tuple_setitem(tup, i, shape[i])
|
|
|
|
# numpy checks if the input arrays have the same shape as `shape`
|
|
outs = []
|
|
for array in literal_unroll(args):
|
|
outs.append(np.broadcast_to(np.asarray(array), tup))
|
|
return outs
|
|
|
|
return impl
|
|
|
|
|
|
def fancy_setslice(context, builder, sig, args, index_types, indices):
|
|
"""
|
|
Implement slice assignment for arrays. This implementation works for
|
|
basic as well as fancy indexing, since there's no functional difference
|
|
between the two for indexed assignment.
|
|
"""
|
|
aryty, _, srcty = sig.args
|
|
ary, _, src = args
|
|
|
|
ary = make_array(aryty)(context, builder, ary)
|
|
dest_shapes = cgutils.unpack_tuple(builder, ary.shape)
|
|
dest_strides = cgutils.unpack_tuple(builder, ary.strides)
|
|
dest_data = ary.data
|
|
|
|
indexer = FancyIndexer(context, builder, aryty, ary,
|
|
index_types, indices)
|
|
indexer.prepare()
|
|
|
|
if isinstance(srcty, types.Buffer):
|
|
# Source is an array
|
|
src_dtype = srcty.dtype
|
|
index_shape = indexer.get_shape()
|
|
src = make_array(srcty)(context, builder, src)
|
|
# Broadcast source array to shape
|
|
srcty, src = _broadcast_to_shape(context, builder, srcty, src,
|
|
index_shape)
|
|
src_shapes = cgutils.unpack_tuple(builder, src.shape)
|
|
src_strides = cgutils.unpack_tuple(builder, src.strides)
|
|
src_data = src.data
|
|
|
|
# Check shapes are equal
|
|
shape_error = cgutils.false_bit
|
|
assert len(index_shape) == len(src_shapes)
|
|
|
|
for u, v in zip(src_shapes, index_shape):
|
|
shape_error = builder.or_(shape_error,
|
|
builder.icmp_signed('!=', u, v))
|
|
|
|
with builder.if_then(shape_error, likely=False):
|
|
msg = "cannot assign slice from input of different size"
|
|
context.call_conv.return_user_exc(builder, ValueError, (msg,))
|
|
|
|
# Check for array overlap
|
|
src_start, src_end = get_array_memory_extents(context, builder, srcty,
|
|
src, src_shapes,
|
|
src_strides, src_data)
|
|
|
|
dest_lower, dest_upper = indexer.get_offset_bounds(dest_strides,
|
|
ary.itemsize)
|
|
dest_start, dest_end = compute_memory_extents(context, builder,
|
|
dest_lower, dest_upper,
|
|
dest_data)
|
|
|
|
use_copy = extents_may_overlap(context, builder, src_start, src_end,
|
|
dest_start, dest_end)
|
|
|
|
src_getitem, src_cleanup = maybe_copy_source(context, builder, use_copy,
|
|
srcty, src, src_shapes,
|
|
src_strides, src_data)
|
|
|
|
elif isinstance(srcty, types.Sequence):
|
|
src_dtype = srcty.dtype
|
|
|
|
# Check shape is equal to sequence length
|
|
index_shape = indexer.get_shape()
|
|
assert len(index_shape) == 1
|
|
len_impl = context.get_function(len, signature(types.intp, srcty))
|
|
seq_len = len_impl(builder, (src,))
|
|
|
|
shape_error = builder.icmp_signed('!=', index_shape[0], seq_len)
|
|
|
|
with builder.if_then(shape_error, likely=False):
|
|
msg = "cannot assign slice from input of different size"
|
|
context.call_conv.return_user_exc(builder, ValueError, (msg,))
|
|
|
|
def src_getitem(source_indices):
|
|
idx, = source_indices
|
|
getitem_impl = context.get_function(
|
|
operator.getitem,
|
|
signature(src_dtype, srcty, types.intp),
|
|
)
|
|
return getitem_impl(builder, (src, idx))
|
|
|
|
def src_cleanup():
|
|
pass
|
|
|
|
else:
|
|
# Source is a scalar (broadcast or not, depending on destination
|
|
# shape).
|
|
src_dtype = srcty
|
|
|
|
def src_getitem(source_indices):
|
|
return src
|
|
|
|
def src_cleanup():
|
|
pass
|
|
|
|
zero = context.get_constant(types.uintp, 0)
|
|
# Loop on destination and copy from source to destination
|
|
dest_indices, counts = indexer.begin_loops()
|
|
|
|
# Source is iterated in natural order
|
|
|
|
# Counts represent a counter for the number of times a specified axis
|
|
# is being accessed, during setitem they are used as source
|
|
# indices
|
|
counts = list(counts)
|
|
|
|
# We need to artifically introduce the index zero wherever a
|
|
# newaxis is present within the indexer. These always remain
|
|
# zero.
|
|
for i in indexer.newaxes:
|
|
counts.insert(i, zero)
|
|
|
|
source_indices = [c for c in counts if c is not None]
|
|
|
|
val = src_getitem(source_indices)
|
|
|
|
# Cast to the destination dtype (cross-dtype slice assignment is allowed)
|
|
val = context.cast(builder, val, src_dtype, aryty.dtype)
|
|
|
|
# No need to check for wraparound, as the indexers all ensure
|
|
# a positive index is returned.
|
|
dest_ptr = cgutils.get_item_pointer2(context, builder, dest_data,
|
|
dest_shapes, dest_strides,
|
|
aryty.layout, dest_indices,
|
|
wraparound=False,
|
|
boundscheck=context.enable_boundscheck)
|
|
store_item(context, builder, aryty, val, dest_ptr)
|
|
|
|
indexer.end_loops()
|
|
|
|
src_cleanup()
|
|
|
|
return context.get_dummy_value()
|
|
|
|
|
|
# ------------------------------------------------------------------------------
|
|
# Shape / layout altering
|
|
|
|
def vararg_to_tuple(context, builder, sig, args):
|
|
aryty = sig.args[0]
|
|
dimtys = sig.args[1:]
|
|
# values
|
|
ary = args[0]
|
|
dims = args[1:]
|
|
# coerce all types to intp
|
|
dims = [context.cast(builder, val, ty, types.intp)
|
|
for ty, val in zip(dimtys, dims)]
|
|
# make a tuple
|
|
shape = cgutils.pack_array(builder, dims, dims[0].type)
|
|
|
|
shapety = types.UniTuple(dtype=types.intp, count=len(dims))
|
|
new_sig = typing.signature(sig.return_type, aryty, shapety)
|
|
new_args = ary, shape
|
|
|
|
return new_sig, new_args
|
|
|
|
|
|
@lower_builtin('array.transpose', types.Array)
|
|
def array_transpose(context, builder, sig, args):
|
|
return array_T(context, builder, sig.args[0], args[0])
|
|
|
|
|
|
def permute_arrays(axis, shape, strides):
|
|
if len(axis) != len(set(axis)):
|
|
raise ValueError("repeated axis in transpose")
|
|
dim = len(shape)
|
|
for x in axis:
|
|
if x >= dim or abs(x) > dim:
|
|
raise ValueError("axis is out of bounds for array of "
|
|
"given dimension")
|
|
|
|
shape[:] = shape[axis]
|
|
strides[:] = strides[axis]
|
|
|
|
|
|
# Transposing an array involves permuting the shape and strides of the array
|
|
# based on the given axes.
|
|
@lower_builtin('array.transpose', types.Array, types.BaseTuple)
|
|
def array_transpose_tuple(context, builder, sig, args):
|
|
aryty = sig.args[0]
|
|
ary = make_array(aryty)(context, builder, args[0])
|
|
|
|
axisty, axis = sig.args[1], args[1]
|
|
num_axis, dtype = axisty.count, axisty.dtype
|
|
|
|
ll_intp = context.get_value_type(types.intp)
|
|
ll_ary_size = ir.ArrayType(ll_intp, num_axis)
|
|
|
|
# Allocate memory for axes, shapes, and strides arrays.
|
|
arys = [axis, ary.shape, ary.strides]
|
|
ll_arys = [cgutils.alloca_once(builder, ll_ary_size) for _ in arys]
|
|
|
|
# Store axes, shapes, and strides arrays to the allocated memory.
|
|
for src, dst in zip(arys, ll_arys):
|
|
builder.store(src, dst)
|
|
|
|
np_ary_ty = types.Array(dtype=dtype, ndim=1, layout='C')
|
|
np_itemsize = context.get_constant(types.intp,
|
|
context.get_abi_sizeof(ll_intp))
|
|
|
|
# Form NumPy arrays for axes, shapes, and strides arrays.
|
|
np_arys = [make_array(np_ary_ty)(context, builder) for _ in arys]
|
|
|
|
# Roughly, `np_ary = np.array(ll_ary)` for each of axes, shapes, and strides
|
|
for np_ary, ll_ary in zip(np_arys, ll_arys):
|
|
populate_array(np_ary,
|
|
data=builder.bitcast(ll_ary, ll_intp.as_pointer()),
|
|
shape=[context.get_constant(types.intp, num_axis)],
|
|
strides=[np_itemsize],
|
|
itemsize=np_itemsize,
|
|
meminfo=None)
|
|
|
|
# Pass NumPy arrays formed above to permute_arrays function that permutes
|
|
# shapes and strides based on axis contents.
|
|
context.compile_internal(builder, permute_arrays,
|
|
typing.signature(types.void,
|
|
np_ary_ty, np_ary_ty, np_ary_ty),
|
|
[a._getvalue() for a in np_arys])
|
|
|
|
# Make a new array based on permuted shape and strides and return it.
|
|
ret = make_array(sig.return_type)(context, builder)
|
|
populate_array(ret,
|
|
data=ary.data,
|
|
shape=builder.load(ll_arys[1]),
|
|
strides=builder.load(ll_arys[2]),
|
|
itemsize=ary.itemsize,
|
|
meminfo=ary.meminfo,
|
|
parent=ary.parent)
|
|
res = ret._getvalue()
|
|
return impl_ret_borrowed(context, builder, sig.return_type, res)
|
|
|
|
|
|
@lower_builtin('array.transpose', types.Array, types.VarArg(types.Any))
|
|
def array_transpose_vararg(context, builder, sig, args):
|
|
new_sig, new_args = vararg_to_tuple(context, builder, sig, args)
|
|
return array_transpose_tuple(context, builder, new_sig, new_args)
|
|
|
|
|
|
@overload(np.transpose)
|
|
def numpy_transpose(a, axes=None):
|
|
if isinstance(a, types.BaseTuple):
|
|
raise errors.UnsupportedError("np.transpose does not accept tuples")
|
|
|
|
if axes is None:
|
|
def np_transpose_impl(a, axes=None):
|
|
return a.transpose()
|
|
else:
|
|
def np_transpose_impl(a, axes=None):
|
|
return a.transpose(axes)
|
|
|
|
return np_transpose_impl
|
|
|
|
|
|
@lower_getattr(types.Array, 'T')
|
|
def array_T(context, builder, typ, value):
|
|
if typ.ndim <= 1:
|
|
res = value
|
|
else:
|
|
ary = make_array(typ)(context, builder, value)
|
|
ret = make_array(typ)(context, builder)
|
|
shapes = cgutils.unpack_tuple(builder, ary.shape, typ.ndim)
|
|
strides = cgutils.unpack_tuple(builder, ary.strides, typ.ndim)
|
|
populate_array(ret,
|
|
data=ary.data,
|
|
shape=cgutils.pack_array(builder, shapes[::-1]),
|
|
strides=cgutils.pack_array(builder, strides[::-1]),
|
|
itemsize=ary.itemsize,
|
|
meminfo=ary.meminfo,
|
|
parent=ary.parent)
|
|
res = ret._getvalue()
|
|
return impl_ret_borrowed(context, builder, typ, res)
|
|
|
|
|
|
@overload(np.logspace)
|
|
def numpy_logspace(start, stop, num=50):
|
|
if not isinstance(start, types.Number):
|
|
raise errors.TypingError('The first argument "start" must be a number')
|
|
if not isinstance(stop, types.Number):
|
|
raise errors.TypingError('The second argument "stop" must be a number')
|
|
if not isinstance(num, (int, types.Integer)):
|
|
raise errors.TypingError('The third argument "num" must be an integer')
|
|
|
|
def impl(start, stop, num=50):
|
|
y = np.linspace(start, stop, num)
|
|
return np.power(10.0, y)
|
|
|
|
return impl
|
|
|
|
|
|
@overload(np.geomspace)
|
|
def numpy_geomspace(start, stop, num=50):
|
|
if not isinstance(start, types.Number):
|
|
msg = 'The argument "start" must be a number'
|
|
raise errors.TypingError(msg)
|
|
|
|
if not isinstance(stop, types.Number):
|
|
msg = 'The argument "stop" must be a number'
|
|
raise errors.TypingError(msg)
|
|
|
|
if not isinstance(num, (int, types.Integer)):
|
|
msg = 'The argument "num" must be an integer'
|
|
raise errors.TypingError(msg)
|
|
|
|
if any(isinstance(arg, types.Complex) for arg in [start, stop]):
|
|
result_dtype = from_dtype(np.result_type(as_dtype(start),
|
|
as_dtype(stop), None))
|
|
|
|
def impl(start, stop, num=50):
|
|
if start == 0 or stop == 0:
|
|
raise ValueError('Geometric sequence cannot include zero')
|
|
start = result_dtype(start)
|
|
stop = result_dtype(stop)
|
|
both_imaginary = (start.real == 0) & (stop.real == 0)
|
|
both_negative = (np.sign(start) == -1) & (np.sign(stop) == -1)
|
|
out_sign = 1
|
|
if both_imaginary:
|
|
start = start.imag
|
|
stop = stop.imag
|
|
out_sign = 1j
|
|
if both_negative:
|
|
start = -start
|
|
stop = -stop
|
|
out_sign = -out_sign
|
|
logstart = np.log10(start)
|
|
logstop = np.log10(stop)
|
|
result = np.logspace(logstart, logstop, num)
|
|
# Make sure the endpoints match the start and stop arguments.
|
|
# This is necessary because np.exp(np.log(x)) is not necessarily
|
|
# equal to x.
|
|
if num > 0:
|
|
result[0] = start
|
|
if num > 1:
|
|
result[-1] = stop
|
|
return out_sign * result
|
|
|
|
else:
|
|
def impl(start, stop, num=50):
|
|
if start == 0 or stop == 0:
|
|
raise ValueError('Geometric sequence cannot include zero')
|
|
both_negative = (np.sign(start) == -1) & (np.sign(stop) == -1)
|
|
out_sign = 1
|
|
if both_negative:
|
|
start = -start
|
|
stop = -stop
|
|
out_sign = -out_sign
|
|
logstart = np.log10(start)
|
|
logstop = np.log10(stop)
|
|
result = np.logspace(logstart, logstop, num)
|
|
# Make sure the endpoints match the start and stop arguments.
|
|
# This is necessary because np.exp(np.log(x)) is not necessarily
|
|
# equal to x.
|
|
if num > 0:
|
|
result[0] = start
|
|
if num > 1:
|
|
result[-1] = stop
|
|
return out_sign * result
|
|
|
|
return impl
|
|
|
|
|
|
@overload(np.rot90)
|
|
def numpy_rot90(m, k=1):
|
|
# supporting axes argument it needs to be included in np.flip
|
|
if not isinstance(k, (int, types.Integer)):
|
|
raise errors.TypingError('The second argument "k" must be an integer')
|
|
if not isinstance(m, types.Array):
|
|
raise errors.TypingError('The first argument "m" must be an array')
|
|
|
|
if m.ndim < 2:
|
|
raise errors.NumbaValueError('Input must be >= 2-d.')
|
|
|
|
def impl(m, k=1):
|
|
k = k % 4
|
|
if k == 0:
|
|
return m[:]
|
|
elif k == 1:
|
|
return np.swapaxes(np.fliplr(m), 0, 1)
|
|
elif k == 2:
|
|
return np.flipud(np.fliplr(m))
|
|
elif k == 3:
|
|
return np.fliplr(np.swapaxes(m, 0, 1))
|
|
else:
|
|
raise AssertionError # unreachable
|
|
|
|
return impl
|
|
|
|
|
|
def _attempt_nocopy_reshape(context, builder, aryty, ary,
|
|
newnd, newshape, newstrides):
|
|
"""
|
|
Call into Numba_attempt_nocopy_reshape() for the given array type
|
|
and instance, and the specified new shape.
|
|
|
|
Return value is non-zero if successful, and the array pointed to
|
|
by *newstrides* will be filled up with the computed results.
|
|
"""
|
|
ll_intp = context.get_value_type(types.intp)
|
|
ll_intp_star = ll_intp.as_pointer()
|
|
ll_intc = context.get_value_type(types.intc)
|
|
fnty = ir.FunctionType(ll_intc, [
|
|
# nd, *dims, *strides
|
|
ll_intp, ll_intp_star, ll_intp_star,
|
|
# newnd, *newdims, *newstrides
|
|
ll_intp, ll_intp_star, ll_intp_star,
|
|
# itemsize, is_f_order
|
|
ll_intp, ll_intc])
|
|
fn = cgutils.get_or_insert_function(builder.module, fnty,
|
|
"numba_attempt_nocopy_reshape")
|
|
|
|
nd = ll_intp(aryty.ndim)
|
|
shape = cgutils.gep_inbounds(builder, ary._get_ptr_by_name('shape'), 0, 0)
|
|
strides = cgutils.gep_inbounds(builder, ary._get_ptr_by_name('strides'),
|
|
0, 0)
|
|
newnd = ll_intp(newnd)
|
|
newshape = cgutils.gep_inbounds(builder, newshape, 0, 0)
|
|
newstrides = cgutils.gep_inbounds(builder, newstrides, 0, 0)
|
|
is_f_order = ll_intc(0)
|
|
res = builder.call(fn, [nd, shape, strides,
|
|
newnd, newshape, newstrides,
|
|
ary.itemsize, is_f_order])
|
|
return res
|
|
|
|
|
|
def normalize_reshape_value(origsize, shape):
|
|
num_neg_value = 0
|
|
known_size = 1
|
|
for ax, s in enumerate(shape):
|
|
if s < 0:
|
|
num_neg_value += 1
|
|
neg_ax = ax
|
|
else:
|
|
known_size *= s
|
|
|
|
if num_neg_value == 0:
|
|
if origsize != known_size:
|
|
raise ValueError("total size of new array must be unchanged")
|
|
|
|
elif num_neg_value == 1:
|
|
# Infer negative dimension
|
|
if known_size == 0:
|
|
inferred = 0
|
|
ok = origsize == 0
|
|
else:
|
|
inferred = origsize // known_size
|
|
ok = origsize % known_size == 0
|
|
if not ok:
|
|
raise ValueError("total size of new array must be unchanged")
|
|
shape[neg_ax] = inferred
|
|
|
|
else:
|
|
raise ValueError("multiple negative shape values")
|
|
|
|
|
|
@lower_builtin('array.reshape', types.Array, types.BaseTuple)
|
|
def array_reshape(context, builder, sig, args):
|
|
aryty = sig.args[0]
|
|
retty = sig.return_type
|
|
|
|
shapety = sig.args[1]
|
|
shape = args[1]
|
|
|
|
ll_intp = context.get_value_type(types.intp)
|
|
ll_shape = ir.ArrayType(ll_intp, shapety.count)
|
|
|
|
ary = make_array(aryty)(context, builder, args[0])
|
|
|
|
# We will change the target shape in this slot
|
|
# (see normalize_reshape_value() below)
|
|
newshape = cgutils.alloca_once(builder, ll_shape)
|
|
builder.store(shape, newshape)
|
|
|
|
# Create a shape array pointing to the value of newshape.
|
|
# (roughly, `shape_ary = np.array(ary.shape)`)
|
|
shape_ary_ty = types.Array(dtype=shapety.dtype, ndim=1, layout='C')
|
|
shape_ary = make_array(shape_ary_ty)(context, builder)
|
|
shape_itemsize = context.get_constant(types.intp,
|
|
context.get_abi_sizeof(ll_intp))
|
|
populate_array(shape_ary,
|
|
data=builder.bitcast(newshape, ll_intp.as_pointer()),
|
|
shape=[context.get_constant(types.intp, shapety.count)],
|
|
strides=[shape_itemsize],
|
|
itemsize=shape_itemsize,
|
|
meminfo=None)
|
|
|
|
# Compute the original array size
|
|
size = ary.nitems
|
|
|
|
# Call our normalizer which will fix the shape array in case of negative
|
|
# shape value
|
|
context.compile_internal(builder, normalize_reshape_value,
|
|
typing.signature(types.void,
|
|
types.uintp, shape_ary_ty),
|
|
[size, shape_ary._getvalue()])
|
|
|
|
# Perform reshape (nocopy)
|
|
newnd = shapety.count
|
|
newstrides = cgutils.alloca_once(builder, ll_shape)
|
|
|
|
ok = _attempt_nocopy_reshape(context, builder, aryty, ary, newnd,
|
|
newshape, newstrides)
|
|
fail = builder.icmp_unsigned('==', ok, ok.type(0))
|
|
|
|
with builder.if_then(fail):
|
|
msg = "incompatible shape for array"
|
|
context.call_conv.return_user_exc(builder, NotImplementedError, (msg,))
|
|
|
|
ret = make_array(retty)(context, builder)
|
|
populate_array(ret,
|
|
data=ary.data,
|
|
shape=builder.load(newshape),
|
|
strides=builder.load(newstrides),
|
|
itemsize=ary.itemsize,
|
|
meminfo=ary.meminfo,
|
|
parent=ary.parent)
|
|
res = ret._getvalue()
|
|
return impl_ret_borrowed(context, builder, sig.return_type, res)
|
|
|
|
|
|
@lower_builtin('array.reshape', types.Array, types.VarArg(types.Any))
|
|
def array_reshape_vararg(context, builder, sig, args):
|
|
new_sig, new_args = vararg_to_tuple(context, builder, sig, args)
|
|
return array_reshape(context, builder, new_sig, new_args)
|
|
|
|
|
|
@overload(np.reshape)
|
|
def np_reshape(a, newshape):
|
|
def np_reshape_impl(a, newshape):
|
|
return a.reshape(newshape)
|
|
return np_reshape_impl
|
|
|
|
|
|
@overload(np.resize)
|
|
def numpy_resize(a, new_shape):
|
|
|
|
if not type_can_asarray(a):
|
|
msg = 'The argument "a" must be array-like'
|
|
raise errors.TypingError(msg)
|
|
|
|
if not ((isinstance(new_shape, types.UniTuple)
|
|
and
|
|
isinstance(new_shape.dtype, types.Integer))
|
|
or
|
|
isinstance(new_shape, types.Integer)):
|
|
msg = ('The argument "new_shape" must be an integer or '
|
|
'a tuple of integers')
|
|
raise errors.TypingError(msg)
|
|
|
|
def impl(a, new_shape):
|
|
a = np.asarray(a)
|
|
a = np.ravel(a)
|
|
|
|
if isinstance(new_shape, tuple):
|
|
new_size = 1
|
|
for dim_length in np.asarray(new_shape):
|
|
new_size *= dim_length
|
|
if dim_length < 0:
|
|
msg = 'All elements of `new_shape` must be non-negative'
|
|
raise ValueError(msg)
|
|
else:
|
|
if new_shape < 0:
|
|
msg2 = 'All elements of `new_shape` must be non-negative'
|
|
raise ValueError(msg2)
|
|
new_size = new_shape
|
|
|
|
if a.size == 0:
|
|
return np.zeros(new_shape).astype(a.dtype)
|
|
|
|
repeats = -(-new_size // a.size) # ceil division
|
|
res = a
|
|
for i in range(repeats - 1):
|
|
res = np.concatenate((res, a))
|
|
res = res[:new_size]
|
|
|
|
return np.reshape(res, new_shape)
|
|
|
|
return impl
|
|
|
|
|
|
@overload(np.append)
|
|
def np_append(arr, values, axis=None):
|
|
|
|
if not type_can_asarray(arr):
|
|
raise errors.TypingError('The first argument "arr" must be array-like')
|
|
|
|
if not type_can_asarray(values):
|
|
raise errors.TypingError('The second argument "values" must be '
|
|
'array-like')
|
|
|
|
if is_nonelike(axis):
|
|
def impl(arr, values, axis=None):
|
|
arr = np.ravel(np.asarray(arr))
|
|
values = np.ravel(np.asarray(values))
|
|
return np.concatenate((arr, values))
|
|
else:
|
|
|
|
if not isinstance(axis, types.Integer):
|
|
raise errors.TypingError('The third argument "axis" must be an '
|
|
'integer')
|
|
|
|
def impl(arr, values, axis=None):
|
|
return np.concatenate((arr, values), axis=axis)
|
|
return impl
|
|
|
|
|
|
@lower_builtin('array.ravel', types.Array)
|
|
def array_ravel(context, builder, sig, args):
|
|
# Only support no argument version (default order='C')
|
|
def imp_nocopy(ary):
|
|
"""No copy version"""
|
|
return ary.reshape(ary.size)
|
|
|
|
def imp_copy(ary):
|
|
"""Copy version"""
|
|
return ary.flatten()
|
|
|
|
# If the input array is C layout already, use the nocopy version
|
|
if sig.args[0].layout == 'C':
|
|
imp = imp_nocopy
|
|
# otherwise, use flatten under-the-hood
|
|
else:
|
|
imp = imp_copy
|
|
|
|
res = context.compile_internal(builder, imp, sig, args)
|
|
res = impl_ret_new_ref(context, builder, sig.return_type, res)
|
|
return res
|
|
|
|
|
|
@lower_builtin(np.ravel, types.Array)
|
|
def np_ravel(context, builder, sig, args):
|
|
def np_ravel_impl(a):
|
|
return a.ravel()
|
|
|
|
return context.compile_internal(builder, np_ravel_impl, sig, args)
|
|
|
|
|
|
@lower_builtin('array.flatten', types.Array)
|
|
def array_flatten(context, builder, sig, args):
|
|
# Only support flattening to C layout currently.
|
|
def imp(ary):
|
|
return ary.copy().reshape(ary.size)
|
|
|
|
res = context.compile_internal(builder, imp, sig, args)
|
|
res = impl_ret_new_ref(context, builder, sig.return_type, res)
|
|
return res
|
|
|
|
|
|
@register_jitable
|
|
def _np_clip_impl(a, a_min, a_max, out):
|
|
# Both a_min and a_max are numpy arrays
|
|
ret = np.empty_like(a) if out is None else out
|
|
a_b, a_min_b, a_max_b = np.broadcast_arrays(a, a_min, a_max)
|
|
for index in np.ndindex(a_b.shape):
|
|
val_a = a_b[index]
|
|
val_a_min = a_min_b[index]
|
|
val_a_max = a_max_b[index]
|
|
ret[index] = min(max(val_a, val_a_min), val_a_max)
|
|
|
|
return ret
|
|
|
|
|
|
@register_jitable
|
|
def _np_clip_impl_none(a, b, use_min, out):
|
|
for index in np.ndindex(a.shape):
|
|
val_a = a[index]
|
|
val_b = b[index]
|
|
if use_min:
|
|
out[index] = min(val_a, val_b)
|
|
else:
|
|
out[index] = max(val_a, val_b)
|
|
return out
|
|
|
|
|
|
@overload(np.clip)
|
|
def np_clip(a, a_min, a_max, out=None):
|
|
if not type_can_asarray(a):
|
|
raise errors.TypingError('The argument "a" must be array-like')
|
|
|
|
if (not isinstance(a_min, types.NoneType) and
|
|
not type_can_asarray(a_min)):
|
|
raise errors.TypingError(('The argument "a_min" must be a number '
|
|
'or an array-like'))
|
|
|
|
if (not isinstance(a_max, types.NoneType) and
|
|
not type_can_asarray(a_max)):
|
|
raise errors.TypingError('The argument "a_max" must be a number '
|
|
'or an array-like')
|
|
|
|
if not (isinstance(out, types.Array) or is_nonelike(out)):
|
|
msg = 'The argument "out" must be an array if it is provided'
|
|
raise errors.TypingError(msg)
|
|
|
|
# TODO: support scalar a (issue #3469)
|
|
a_min_is_none = a_min is None or isinstance(a_min, types.NoneType)
|
|
a_max_is_none = a_max is None or isinstance(a_max, types.NoneType)
|
|
|
|
if a_min_is_none and a_max_is_none:
|
|
# Raises value error when both a_min and a_max are None
|
|
def np_clip_nn(a, a_min, a_max, out=None):
|
|
raise ValueError("array_clip: must set either max or min")
|
|
|
|
return np_clip_nn
|
|
|
|
a_min_is_scalar = isinstance(a_min, types.Number)
|
|
a_max_is_scalar = isinstance(a_max, types.Number)
|
|
|
|
if a_min_is_scalar and a_max_is_scalar:
|
|
def np_clip_ss(a, a_min, a_max, out=None):
|
|
# a_min and a_max are scalars
|
|
# since their shape will be empty
|
|
# so broadcasting is not needed at all
|
|
ret = np.empty_like(a) if out is None else out
|
|
for index in np.ndindex(a.shape):
|
|
val_a = a[index]
|
|
ret[index] = min(max(val_a, a_min), a_max)
|
|
|
|
return ret
|
|
|
|
return np_clip_ss
|
|
elif a_min_is_scalar and not a_max_is_scalar:
|
|
if a_max_is_none:
|
|
def np_clip_sn(a, a_min, a_max, out=None):
|
|
# a_min is a scalar
|
|
# since its shape will be empty
|
|
# so broadcasting is not needed at all
|
|
ret = np.empty_like(a) if out is None else out
|
|
for index in np.ndindex(a.shape):
|
|
val_a = a[index]
|
|
ret[index] = max(val_a, a_min)
|
|
|
|
return ret
|
|
|
|
return np_clip_sn
|
|
else:
|
|
def np_clip_sa(a, a_min, a_max, out=None):
|
|
# a_min is a scalar
|
|
# since its shape will be empty
|
|
# broadcast it to shape of a
|
|
# by using np.full_like
|
|
a_min_full = np.full_like(a, a_min)
|
|
return _np_clip_impl(a, a_min_full, a_max, out)
|
|
|
|
return np_clip_sa
|
|
elif not a_min_is_scalar and a_max_is_scalar:
|
|
if a_min_is_none:
|
|
def np_clip_ns(a, a_min, a_max, out=None):
|
|
# a_max is a scalar
|
|
# since its shape will be empty
|
|
# so broadcasting is not needed at all
|
|
ret = np.empty_like(a) if out is None else out
|
|
for index in np.ndindex(a.shape):
|
|
val_a = a[index]
|
|
ret[index] = min(val_a, a_max)
|
|
|
|
return ret
|
|
|
|
return np_clip_ns
|
|
else:
|
|
def np_clip_as(a, a_min, a_max, out=None):
|
|
# a_max is a scalar
|
|
# since its shape will be empty
|
|
# broadcast it to shape of a
|
|
# by using np.full_like
|
|
a_max_full = np.full_like(a, a_max)
|
|
return _np_clip_impl(a, a_min, a_max_full, out)
|
|
|
|
return np_clip_as
|
|
else:
|
|
# Case where exactly one of a_min or a_max is None
|
|
if a_min_is_none:
|
|
def np_clip_na(a, a_min, a_max, out=None):
|
|
# a_max is a numpy array but a_min is None
|
|
ret = np.empty_like(a) if out is None else out
|
|
a_b, a_max_b = np.broadcast_arrays(a, a_max)
|
|
return _np_clip_impl_none(a_b, a_max_b, True, ret)
|
|
|
|
return np_clip_na
|
|
elif a_max_is_none:
|
|
def np_clip_an(a, a_min, a_max, out=None):
|
|
# a_min is a numpy array but a_max is None
|
|
ret = np.empty_like(a) if out is None else out
|
|
a_b, a_min_b = np.broadcast_arrays(a, a_min)
|
|
return _np_clip_impl_none(a_b, a_min_b, False, ret)
|
|
|
|
return np_clip_an
|
|
else:
|
|
def np_clip_aa(a, a_min, a_max, out=None):
|
|
# Both a_min and a_max are clearly arrays
|
|
# because none of the above branches
|
|
# returned
|
|
return _np_clip_impl(a, a_min, a_max, out)
|
|
|
|
return np_clip_aa
|
|
|
|
|
|
@overload_method(types.Array, 'clip')
|
|
def array_clip(a, a_min=None, a_max=None, out=None):
|
|
def impl(a, a_min=None, a_max=None, out=None):
|
|
return np.clip(a, a_min, a_max, out)
|
|
return impl
|
|
|
|
|
|
def _change_dtype(context, builder, oldty, newty, ary):
|
|
"""
|
|
Attempt to fix up *ary* for switching from *oldty* to *newty*.
|
|
|
|
See Numpy's array_descr_set()
|
|
(np/core/src/multiarray/getset.c).
|
|
Attempt to fix the array's shape and strides for a new dtype.
|
|
False is returned on failure, True on success.
|
|
"""
|
|
assert oldty.ndim == newty.ndim
|
|
assert oldty.layout == newty.layout
|
|
|
|
new_layout = ord(newty.layout)
|
|
any_layout = ord('A')
|
|
c_layout = ord('C')
|
|
f_layout = ord('F')
|
|
|
|
int8 = types.int8
|
|
|
|
def imp(nd, dims, strides, old_itemsize, new_itemsize, layout):
|
|
# Attempt to update the layout due to limitation of the numba
|
|
# type system.
|
|
if layout == any_layout:
|
|
# Test rightmost stride to be contiguous
|
|
if strides[-1] == old_itemsize:
|
|
# Process this as if it is C contiguous
|
|
layout = int8(c_layout)
|
|
# Test leftmost stride to be F contiguous
|
|
elif strides[0] == old_itemsize:
|
|
# Process this as if it is F contiguous
|
|
layout = int8(f_layout)
|
|
|
|
if old_itemsize != new_itemsize and (layout == any_layout or nd == 0):
|
|
return False
|
|
|
|
if layout == c_layout:
|
|
i = nd - 1
|
|
else:
|
|
i = 0
|
|
|
|
if new_itemsize < old_itemsize:
|
|
# If it is compatible, increase the size of the dimension
|
|
# at the end (or at the front if F-contiguous)
|
|
if (old_itemsize % new_itemsize) != 0:
|
|
return False
|
|
|
|
newdim = old_itemsize // new_itemsize
|
|
dims[i] *= newdim
|
|
strides[i] = new_itemsize
|
|
|
|
elif new_itemsize > old_itemsize:
|
|
# Determine if last (or first if F-contiguous) dimension
|
|
# is compatible
|
|
bytelength = dims[i] * old_itemsize
|
|
if (bytelength % new_itemsize) != 0:
|
|
return False
|
|
|
|
dims[i] = bytelength // new_itemsize
|
|
strides[i] = new_itemsize
|
|
|
|
else:
|
|
# Same item size: nothing to do (this also works for
|
|
# non-contiguous arrays).
|
|
pass
|
|
|
|
return True
|
|
|
|
old_itemsize = context.get_constant(types.intp,
|
|
get_itemsize(context, oldty))
|
|
new_itemsize = context.get_constant(types.intp,
|
|
get_itemsize(context, newty))
|
|
|
|
nd = context.get_constant(types.intp, newty.ndim)
|
|
shape_data = cgutils.gep_inbounds(builder, ary._get_ptr_by_name('shape'),
|
|
0, 0)
|
|
strides_data = cgutils.gep_inbounds(builder,
|
|
ary._get_ptr_by_name('strides'), 0, 0)
|
|
|
|
shape_strides_array_type = types.Array(dtype=types.intp, ndim=1, layout='C')
|
|
arycls = context.make_array(shape_strides_array_type)
|
|
|
|
shape_constant = cgutils.pack_array(builder,
|
|
[context.get_constant(types.intp,
|
|
newty.ndim)])
|
|
|
|
sizeof_intp = context.get_abi_sizeof(context.get_data_type(types.intp))
|
|
sizeof_intp = context.get_constant(types.intp, sizeof_intp)
|
|
strides_constant = cgutils.pack_array(builder, [sizeof_intp])
|
|
|
|
shape_ary = arycls(context, builder)
|
|
|
|
populate_array(shape_ary,
|
|
data=shape_data,
|
|
shape=shape_constant,
|
|
strides=strides_constant,
|
|
itemsize=sizeof_intp,
|
|
meminfo=None)
|
|
|
|
strides_ary = arycls(context, builder)
|
|
populate_array(strides_ary,
|
|
data=strides_data,
|
|
shape=shape_constant,
|
|
strides=strides_constant,
|
|
itemsize=sizeof_intp,
|
|
meminfo=None)
|
|
|
|
shape = shape_ary._getvalue()
|
|
strides = strides_ary._getvalue()
|
|
args = [nd, shape, strides, old_itemsize, new_itemsize,
|
|
context.get_constant(types.int8, new_layout)]
|
|
|
|
sig = signature(types.boolean,
|
|
types.intp, # nd
|
|
shape_strides_array_type, # dims
|
|
shape_strides_array_type, # strides
|
|
types.intp, # old_itemsize
|
|
types.intp, # new_itemsize
|
|
types.int8, # layout
|
|
)
|
|
|
|
res = context.compile_internal(builder, imp, sig, args)
|
|
update_array_info(newty, ary)
|
|
res = impl_ret_borrowed(context, builder, sig.return_type, res)
|
|
return res
|
|
|
|
|
|
@overload(np.shape)
|
|
def np_shape(a):
|
|
if not type_can_asarray(a):
|
|
raise errors.TypingError("The argument to np.shape must be array-like")
|
|
|
|
def impl(a):
|
|
return np.asarray(a).shape
|
|
return impl
|
|
|
|
# ------------------------------------------------------------------------------
|
|
|
|
|
|
@overload(np.unique)
|
|
def np_unique(ar):
|
|
def np_unique_impl(ar):
|
|
b = np.sort(ar.ravel())
|
|
head = list(b[:1])
|
|
tail = [x for i, x in enumerate(b[1:]) if b[i] != x]
|
|
return np.array(head + tail)
|
|
return np_unique_impl
|
|
|
|
|
|
@overload(np.repeat)
|
|
def np_repeat(a, repeats):
|
|
# Implementation for repeats being a scalar is a module global function
|
|
# (see below) because it might be called from the implementation below.
|
|
|
|
def np_repeat_impl_repeats_array_like(a, repeats):
|
|
# implementation if repeats is an array like
|
|
repeats_array = np.asarray(repeats, dtype=np.int64)
|
|
# if it is a singleton array, invoke the scalar implementation
|
|
if repeats_array.shape[0] == 1:
|
|
return np_repeat_impl_repeats_scaler(a, repeats_array[0])
|
|
if np.any(repeats_array < 0):
|
|
raise ValueError("negative dimensions are not allowed")
|
|
asa = np.asarray(a)
|
|
aravel = asa.ravel()
|
|
n = aravel.shape[0]
|
|
if aravel.shape != repeats_array.shape:
|
|
raise ValueError(
|
|
"operands could not be broadcast together")
|
|
to_return = np.empty(np.sum(repeats_array), dtype=asa.dtype)
|
|
pos = 0
|
|
for i in range(n):
|
|
to_return[pos : pos + repeats_array[i]] = aravel[i]
|
|
pos += repeats_array[i]
|
|
return to_return
|
|
|
|
# type checking
|
|
if isinstance(a, (types.Array,
|
|
types.List,
|
|
types.BaseTuple,
|
|
types.Number,
|
|
types.Boolean,
|
|
)
|
|
):
|
|
if isinstance(repeats, types.Integer):
|
|
return np_repeat_impl_repeats_scaler
|
|
elif isinstance(repeats, (types.Array, types.List)):
|
|
if isinstance(repeats.dtype, types.Integer):
|
|
return np_repeat_impl_repeats_array_like
|
|
|
|
raise errors.TypingError(
|
|
"The repeats argument must be an integer "
|
|
"or an array-like of integer dtype")
|
|
|
|
|
|
@register_jitable
|
|
def np_repeat_impl_repeats_scaler(a, repeats):
|
|
if repeats < 0:
|
|
raise ValueError("negative dimensions are not allowed")
|
|
asa = np.asarray(a)
|
|
aravel = asa.ravel()
|
|
n = aravel.shape[0]
|
|
if repeats == 0:
|
|
return np.empty(0, dtype=asa.dtype)
|
|
elif repeats == 1:
|
|
return np.copy(aravel)
|
|
else:
|
|
to_return = np.empty(n * repeats, dtype=asa.dtype)
|
|
for i in range(n):
|
|
to_return[i * repeats : (i + 1) * repeats] = aravel[i]
|
|
return to_return
|
|
|
|
|
|
@extending.overload_method(types.Array, 'repeat')
|
|
def array_repeat(a, repeats):
|
|
def array_repeat_impl(a, repeats):
|
|
return np.repeat(a, repeats)
|
|
|
|
return array_repeat_impl
|
|
|
|
|
|
@intrinsic
|
|
def _intrin_get_itemsize(tyctx, dtype):
|
|
"""Computes the itemsize of the dtype"""
|
|
sig = types.intp(dtype)
|
|
|
|
def codegen(cgctx, builder, sig, llargs):
|
|
llty = cgctx.get_data_type(sig.args[0].dtype)
|
|
llintp = cgctx.get_data_type(sig.return_type)
|
|
return llintp(cgctx.get_abi_sizeof(llty))
|
|
return sig, codegen
|
|
|
|
|
|
def _compatible_view(a, dtype):
|
|
pass
|
|
|
|
|
|
@overload(_compatible_view, target='generic')
|
|
def ol_compatible_view(a, dtype):
|
|
"""Determines if the array and dtype are compatible for forming a view."""
|
|
# NOTE: NumPy 1.23+ uses this check.
|
|
# Code based on:
|
|
# https://github.com/numpy/numpy/blob/750ad21258cfc00663586d5a466e24f91b48edc7/numpy/core/src/multiarray/getset.c#L500-L555 # noqa: E501
|
|
def impl(a, dtype):
|
|
dtype_size = _intrin_get_itemsize(dtype)
|
|
if dtype_size != a.itemsize:
|
|
# catch forbidden cases
|
|
if a.ndim == 0:
|
|
msg1 = ("Changing the dtype of a 0d array is only supported "
|
|
"if the itemsize is unchanged")
|
|
raise ValueError(msg1)
|
|
else:
|
|
# NumPy has a check here for subarray type conversion which
|
|
# Numba doesn't support
|
|
pass
|
|
|
|
# Resize on last axis only
|
|
axis = a.ndim - 1
|
|
p1 = a.shape[axis] != 1
|
|
p2 = a.size != 0
|
|
p3 = a.strides[axis] != a.itemsize
|
|
if (p1 and p2 and p3):
|
|
msg2 = ("To change to a dtype of a different size, the last "
|
|
"axis must be contiguous")
|
|
raise ValueError(msg2)
|
|
|
|
if dtype_size < a.itemsize:
|
|
if dtype_size == 0 or a.itemsize % dtype_size != 0:
|
|
msg3 = ("When changing to a smaller dtype, its size must "
|
|
"be a divisor of the size of original dtype")
|
|
raise ValueError(msg3)
|
|
else:
|
|
newdim = a.shape[axis] * a.itemsize
|
|
if newdim % dtype_size != 0:
|
|
msg4 = ("When changing to a larger dtype, its size must be "
|
|
"a divisor of the total size in bytes of the last "
|
|
"axis of the array.")
|
|
raise ValueError(msg4)
|
|
return impl
|
|
|
|
|
|
@lower_builtin('array.view', types.Array, types.DTypeSpec)
|
|
def array_view(context, builder, sig, args):
|
|
aryty = sig.args[0]
|
|
retty = sig.return_type
|
|
|
|
ary = make_array(aryty)(context, builder, args[0])
|
|
ret = make_array(retty)(context, builder)
|
|
# Copy all fields, casting the "data" pointer appropriately
|
|
fields = set(ret._datamodel._fields)
|
|
for k in sorted(fields):
|
|
val = getattr(ary, k)
|
|
if k == 'data':
|
|
ptrty = ret.data.type
|
|
ret.data = builder.bitcast(val, ptrty)
|
|
else:
|
|
setattr(ret, k, val)
|
|
|
|
if numpy_version >= (1, 23):
|
|
# NumPy 1.23+ bans views using a dtype that is a different size to that
|
|
# of the array when the last axis is not contiguous. For example, this
|
|
# manifests at runtime when a dtype size altering view is requested
|
|
# on a Fortran ordered array.
|
|
|
|
tyctx = context.typing_context
|
|
fnty = tyctx.resolve_value_type(_compatible_view)
|
|
_compatible_view_sig = fnty.get_call_type(tyctx, (*sig.args,), {})
|
|
impl = context.get_function(fnty, _compatible_view_sig)
|
|
impl(builder, args)
|
|
|
|
ok = _change_dtype(context, builder, aryty, retty, ret)
|
|
fail = builder.icmp_unsigned('==', ok, Constant(ok.type, 0))
|
|
|
|
with builder.if_then(fail):
|
|
msg = "new type not compatible with array"
|
|
context.call_conv.return_user_exc(builder, ValueError, (msg,))
|
|
|
|
res = ret._getvalue()
|
|
return impl_ret_borrowed(context, builder, sig.return_type, res)
|
|
|
|
|
|
# ------------------------------------------------------------------------------
|
|
# Array attributes
|
|
|
|
@lower_getattr(types.Array, "dtype")
|
|
def array_dtype(context, builder, typ, value):
|
|
res = context.get_dummy_value()
|
|
return impl_ret_untracked(context, builder, typ, res)
|
|
|
|
|
|
@lower_getattr(types.Array, "shape")
|
|
@lower_getattr(types.MemoryView, "shape")
|
|
def array_shape(context, builder, typ, value):
|
|
arrayty = make_array(typ)
|
|
array = arrayty(context, builder, value)
|
|
res = array.shape
|
|
return impl_ret_untracked(context, builder, typ, res)
|
|
|
|
|
|
@lower_getattr(types.Array, "strides")
|
|
@lower_getattr(types.MemoryView, "strides")
|
|
def array_strides(context, builder, typ, value):
|
|
arrayty = make_array(typ)
|
|
array = arrayty(context, builder, value)
|
|
res = array.strides
|
|
return impl_ret_untracked(context, builder, typ, res)
|
|
|
|
|
|
@lower_getattr(types.Array, "ndim")
|
|
@lower_getattr(types.MemoryView, "ndim")
|
|
def array_ndim(context, builder, typ, value):
|
|
res = context.get_constant(types.intp, typ.ndim)
|
|
return impl_ret_untracked(context, builder, typ, res)
|
|
|
|
|
|
@lower_getattr(types.Array, "size")
|
|
def array_size(context, builder, typ, value):
|
|
arrayty = make_array(typ)
|
|
array = arrayty(context, builder, value)
|
|
res = array.nitems
|
|
return impl_ret_untracked(context, builder, typ, res)
|
|
|
|
|
|
@lower_getattr(types.Array, "itemsize")
|
|
@lower_getattr(types.MemoryView, "itemsize")
|
|
def array_itemsize(context, builder, typ, value):
|
|
arrayty = make_array(typ)
|
|
array = arrayty(context, builder, value)
|
|
res = array.itemsize
|
|
return impl_ret_untracked(context, builder, typ, res)
|
|
|
|
|
|
@lower_getattr(types.Array, "nbytes")
|
|
@lower_getattr(types.MemoryView, "nbytes")
|
|
def array_nbytes(context, builder, typ, value):
|
|
"""
|
|
nbytes = size * itemsize
|
|
"""
|
|
arrayty = make_array(typ)
|
|
array = arrayty(context, builder, value)
|
|
res = builder.mul(array.nitems, array.itemsize)
|
|
return impl_ret_untracked(context, builder, typ, res)
|
|
|
|
|
|
@lower_getattr(types.MemoryView, "contiguous")
|
|
def array_contiguous(context, builder, typ, value):
|
|
res = context.get_constant(types.boolean, typ.is_contig)
|
|
return impl_ret_untracked(context, builder, typ, res)
|
|
|
|
|
|
@lower_getattr(types.MemoryView, "c_contiguous")
|
|
def array_c_contiguous(context, builder, typ, value):
|
|
res = context.get_constant(types.boolean, typ.is_c_contig)
|
|
return impl_ret_untracked(context, builder, typ, res)
|
|
|
|
|
|
@lower_getattr(types.MemoryView, "f_contiguous")
|
|
def array_f_contiguous(context, builder, typ, value):
|
|
res = context.get_constant(types.boolean, typ.is_f_contig)
|
|
return impl_ret_untracked(context, builder, typ, res)
|
|
|
|
|
|
@lower_getattr(types.MemoryView, "readonly")
|
|
def array_readonly(context, builder, typ, value):
|
|
res = context.get_constant(types.boolean, not typ.mutable)
|
|
return impl_ret_untracked(context, builder, typ, res)
|
|
|
|
|
|
# array.ctypes
|
|
|
|
@lower_getattr(types.Array, "ctypes")
|
|
def array_ctypes(context, builder, typ, value):
|
|
arrayty = make_array(typ)
|
|
array = arrayty(context, builder, value)
|
|
# Create new ArrayCType structure
|
|
act = types.ArrayCTypes(typ)
|
|
ctinfo = context.make_helper(builder, act)
|
|
ctinfo.data = array.data
|
|
ctinfo.meminfo = array.meminfo
|
|
res = ctinfo._getvalue()
|
|
return impl_ret_borrowed(context, builder, act, res)
|
|
|
|
|
|
@lower_getattr(types.ArrayCTypes, "data")
|
|
def array_ctypes_data(context, builder, typ, value):
|
|
ctinfo = context.make_helper(builder, typ, value=value)
|
|
res = ctinfo.data
|
|
# Convert it to an integer
|
|
res = builder.ptrtoint(res, context.get_value_type(types.intp))
|
|
return impl_ret_untracked(context, builder, typ, res)
|
|
|
|
|
|
@lower_cast(types.ArrayCTypes, types.CPointer)
|
|
@lower_cast(types.ArrayCTypes, types.voidptr)
|
|
def array_ctypes_to_pointer(context, builder, fromty, toty, val):
|
|
ctinfo = context.make_helper(builder, fromty, value=val)
|
|
res = ctinfo.data
|
|
res = builder.bitcast(res, context.get_value_type(toty))
|
|
return impl_ret_untracked(context, builder, toty, res)
|
|
|
|
|
|
def _call_contiguous_check(checker, context, builder, aryty, ary):
|
|
"""Helper to invoke the contiguous checker function on an array
|
|
|
|
Args
|
|
----
|
|
checker :
|
|
``numba.numpy_supports.is_contiguous``, or
|
|
``numba.numpy_supports.is_fortran``.
|
|
context : target context
|
|
builder : llvm ir builder
|
|
aryty : numba type
|
|
ary : llvm value
|
|
"""
|
|
ary = make_array(aryty)(context, builder, value=ary)
|
|
tup_intp = types.UniTuple(types.intp, aryty.ndim)
|
|
itemsize = context.get_abi_sizeof(context.get_value_type(aryty.dtype))
|
|
check_sig = signature(types.bool_, tup_intp, tup_intp, types.intp)
|
|
check_args = [ary.shape, ary.strides,
|
|
context.get_constant(types.intp, itemsize)]
|
|
is_contig = context.compile_internal(builder, checker, check_sig,
|
|
check_args)
|
|
return is_contig
|
|
|
|
|
|
# array.flags
|
|
|
|
@lower_getattr(types.Array, "flags")
|
|
def array_flags(context, builder, typ, value):
|
|
flagsobj = context.make_helper(builder, types.ArrayFlags(typ))
|
|
flagsobj.parent = value
|
|
res = flagsobj._getvalue()
|
|
context.nrt.incref(builder, typ, value)
|
|
return impl_ret_new_ref(context, builder, typ, res)
|
|
|
|
|
|
@lower_getattr(types.ArrayFlags, "contiguous")
|
|
@lower_getattr(types.ArrayFlags, "c_contiguous")
|
|
def array_flags_c_contiguous(context, builder, typ, value):
|
|
if typ.array_type.layout != 'C':
|
|
# any layout can still be contiguous
|
|
flagsobj = context.make_helper(builder, typ, value=value)
|
|
res = _call_contiguous_check(is_contiguous, context, builder,
|
|
typ.array_type, flagsobj.parent)
|
|
else:
|
|
val = typ.array_type.layout == 'C'
|
|
res = context.get_constant(types.boolean, val)
|
|
return impl_ret_untracked(context, builder, typ, res)
|
|
|
|
|
|
@lower_getattr(types.ArrayFlags, "f_contiguous")
|
|
def array_flags_f_contiguous(context, builder, typ, value):
|
|
if typ.array_type.layout != 'F':
|
|
# any layout can still be contiguous
|
|
flagsobj = context.make_helper(builder, typ, value=value)
|
|
res = _call_contiguous_check(is_fortran, context, builder,
|
|
typ.array_type, flagsobj.parent)
|
|
else:
|
|
layout = typ.array_type.layout
|
|
val = layout == 'F' if typ.array_type.ndim > 1 else layout in 'CF'
|
|
res = context.get_constant(types.boolean, val)
|
|
return impl_ret_untracked(context, builder, typ, res)
|
|
|
|
|
|
# ------------------------------------------------------------------------------
|
|
# .real / .imag
|
|
|
|
@lower_getattr(types.Array, "real")
|
|
def array_real_part(context, builder, typ, value):
|
|
if typ.dtype in types.complex_domain:
|
|
return array_complex_attr(context, builder, typ, value, attr='real')
|
|
elif typ.dtype in types.number_domain:
|
|
# as an identity function
|
|
return impl_ret_borrowed(context, builder, typ, value)
|
|
else:
|
|
raise NotImplementedError('unsupported .real for {}'.format(type.dtype))
|
|
|
|
|
|
@lower_getattr(types.Array, "imag")
|
|
def array_imag_part(context, builder, typ, value):
|
|
if typ.dtype in types.complex_domain:
|
|
return array_complex_attr(context, builder, typ, value, attr='imag')
|
|
elif typ.dtype in types.number_domain:
|
|
# return a readonly zero array
|
|
sig = signature(typ.copy(readonly=True), typ)
|
|
arrtype, shapes = _parse_empty_like_args(context, builder, sig, [value])
|
|
ary = _empty_nd_impl(context, builder, arrtype, shapes)
|
|
cgutils.memset(builder, ary.data, builder.mul(ary.itemsize,
|
|
ary.nitems), 0)
|
|
return impl_ret_new_ref(context, builder, sig.return_type,
|
|
ary._getvalue())
|
|
else:
|
|
raise NotImplementedError('unsupported .imag for {}'.format(type.dtype))
|
|
|
|
|
|
def array_complex_attr(context, builder, typ, value, attr):
|
|
"""
|
|
Given a complex array, it's memory layout is:
|
|
|
|
R C R C R C
|
|
^ ^ ^
|
|
|
|
(`R` indicates a float for the real part;
|
|
`C` indicates a float for the imaginary part;
|
|
the `^` indicates the start of each element)
|
|
|
|
To get the real part, we can simply change the dtype and itemsize to that
|
|
of the underlying float type. The new layout is:
|
|
|
|
R x R x R x
|
|
^ ^ ^
|
|
|
|
(`x` indicates unused)
|
|
|
|
A load operation will use the dtype to determine the number of bytes to
|
|
load.
|
|
|
|
To get the imaginary part, we shift the pointer by 1 float offset and
|
|
change the dtype and itemsize. The new layout is:
|
|
|
|
x C x C x C
|
|
^ ^ ^
|
|
"""
|
|
if attr not in ['real', 'imag'] or typ.dtype not in types.complex_domain:
|
|
raise NotImplementedError("cannot get attribute `{}`".format(attr))
|
|
|
|
arrayty = make_array(typ)
|
|
array = arrayty(context, builder, value)
|
|
|
|
# sizeof underlying float type
|
|
flty = typ.dtype.underlying_float
|
|
sizeof_flty = context.get_abi_sizeof(context.get_data_type(flty))
|
|
itemsize = array.itemsize.type(sizeof_flty)
|
|
|
|
# cast data pointer to float type
|
|
llfltptrty = context.get_value_type(flty).as_pointer()
|
|
dataptr = builder.bitcast(array.data, llfltptrty)
|
|
|
|
# add offset
|
|
if attr == 'imag':
|
|
dataptr = builder.gep(dataptr, [ir.IntType(32)(1)])
|
|
|
|
# make result
|
|
resultty = typ.copy(dtype=flty, layout='A')
|
|
result = make_array(resultty)(context, builder)
|
|
repl = dict(data=dataptr, itemsize=itemsize)
|
|
cgutils.copy_struct(result, array, repl)
|
|
return impl_ret_borrowed(context, builder, resultty, result._getvalue())
|
|
|
|
|
|
@overload_method(types.Array, 'conj')
|
|
@overload_method(types.Array, 'conjugate')
|
|
def array_conj(arr):
|
|
def impl(arr):
|
|
return np.conj(arr)
|
|
return impl
|
|
|
|
# ------------------------------------------------------------------------------
|
|
# DType attribute
|
|
|
|
|
|
def dtype_type(context, builder, dtypety, dtypeval):
|
|
# Just return a dummy opaque value
|
|
return context.get_dummy_value()
|
|
|
|
|
|
lower_getattr(types.DType, 'type')(dtype_type)
|
|
lower_getattr(types.DType, 'kind')(dtype_type)
|
|
|
|
|
|
# ------------------------------------------------------------------------------
|
|
# static_getitem on Numba numerical types to create "array" types
|
|
|
|
|
|
@lower_builtin('static_getitem', types.NumberClass, types.Any)
|
|
def static_getitem_number_clazz(context, builder, sig, args):
|
|
"""This handles the "static_getitem" when a Numba type is subscripted e.g:
|
|
var = typed.List.empty_list(float64[::1, :])
|
|
It only allows this on simple numerical types. Compound types, like
|
|
records, are not supported.
|
|
"""
|
|
retty = sig.return_type
|
|
if isinstance(retty, types.Array):
|
|
# This isn't used or practically accessible, but has to exist, so just
|
|
# put in a NULL of the right type.
|
|
res = context.get_value_type(retty)(None)
|
|
return impl_ret_untracked(context, builder, retty, res)
|
|
else:
|
|
# This should be unreachable unless the implementation on the Type
|
|
# metaclass is changed.
|
|
msg = ("Unreachable; the definition of __getitem__ on the "
|
|
"numba.types.abstract.Type metaclass should prevent access.")
|
|
raise errors.LoweringError(msg)
|
|
|
|
|
|
# ------------------------------------------------------------------------------
|
|
# Structured / record lookup
|
|
|
|
@lower_getattr_generic(types.Array)
|
|
def array_record_getattr(context, builder, typ, value, attr):
|
|
"""
|
|
Generic getattr() implementation for record arrays: fetch the given
|
|
record member, i.e. a subarray.
|
|
"""
|
|
arrayty = make_array(typ)
|
|
array = arrayty(context, builder, value)
|
|
|
|
rectype = typ.dtype
|
|
if not isinstance(rectype, types.Record):
|
|
raise NotImplementedError("attribute %r of %s not defined"
|
|
% (attr, typ))
|
|
dtype = rectype.typeof(attr)
|
|
offset = rectype.offset(attr)
|
|
|
|
if isinstance(dtype, types.NestedArray):
|
|
resty = typ.copy(
|
|
dtype=dtype.dtype, ndim=typ.ndim + dtype.ndim, layout='A')
|
|
else:
|
|
resty = typ.copy(dtype=dtype, layout='A')
|
|
|
|
raryty = make_array(resty)
|
|
|
|
rary = raryty(context, builder)
|
|
|
|
constoffset = context.get_constant(types.intp, offset)
|
|
|
|
newdataptr = cgutils.pointer_add(
|
|
builder, array.data, constoffset, return_type=rary.data.type,
|
|
)
|
|
if isinstance(dtype, types.NestedArray):
|
|
# new shape = recarray shape + inner dimension from nestedarray
|
|
shape = cgutils.unpack_tuple(builder, array.shape, typ.ndim)
|
|
shape += [context.get_constant(types.intp, i) for i in dtype.shape]
|
|
# new strides = recarray strides + strides of the inner nestedarray
|
|
strides = cgutils.unpack_tuple(builder, array.strides, typ.ndim)
|
|
strides += [context.get_constant(types.intp, i) for i in dtype.strides]
|
|
# New datasize = size of elements of the nestedarray
|
|
datasize = context.get_abi_sizeof(context.get_data_type(dtype.dtype))
|
|
else:
|
|
# New shape, strides, and datasize match the underlying array
|
|
shape = array.shape
|
|
strides = array.strides
|
|
datasize = context.get_abi_sizeof(context.get_data_type(dtype))
|
|
populate_array(rary,
|
|
data=newdataptr,
|
|
shape=shape,
|
|
strides=strides,
|
|
itemsize=context.get_constant(types.intp, datasize),
|
|
meminfo=array.meminfo,
|
|
parent=array.parent)
|
|
res = rary._getvalue()
|
|
return impl_ret_borrowed(context, builder, resty, res)
|
|
|
|
|
|
@lower_builtin('static_getitem', types.Array, types.StringLiteral)
|
|
def array_record_getitem(context, builder, sig, args):
|
|
index = args[1]
|
|
if not isinstance(index, str):
|
|
# This will fallback to normal getitem
|
|
raise NotImplementedError
|
|
return array_record_getattr(context, builder, sig.args[0], args[0], index)
|
|
|
|
|
|
@lower_getattr_generic(types.Record)
|
|
def record_getattr(context, builder, typ, value, attr):
|
|
"""
|
|
Generic getattr() implementation for records: get the given record member.
|
|
"""
|
|
context.sentry_record_alignment(typ, attr)
|
|
offset = typ.offset(attr)
|
|
elemty = typ.typeof(attr)
|
|
|
|
if isinstance(elemty, types.NestedArray):
|
|
# Only a nested array's *data* is stored in a structured array,
|
|
# so we create an array structure to point to that data.
|
|
aryty = make_array(elemty)
|
|
ary = aryty(context, builder)
|
|
dtype = elemty.dtype
|
|
newshape = [context.get_constant(types.intp, s) for s in
|
|
elemty.shape]
|
|
newstrides = [context.get_constant(types.intp, s) for s in
|
|
elemty.strides]
|
|
newdata = cgutils.get_record_member(builder, value, offset,
|
|
context.get_data_type(dtype))
|
|
populate_array(
|
|
ary,
|
|
data=newdata,
|
|
shape=cgutils.pack_array(builder, newshape),
|
|
strides=cgutils.pack_array(builder, newstrides),
|
|
itemsize=context.get_constant(types.intp, elemty.size),
|
|
meminfo=None,
|
|
parent=None,
|
|
)
|
|
res = ary._getvalue()
|
|
return impl_ret_borrowed(context, builder, typ, res)
|
|
else:
|
|
dptr = cgutils.get_record_member(builder, value, offset,
|
|
context.get_data_type(elemty))
|
|
align = None if typ.aligned else 1
|
|
res = context.unpack_value(builder, elemty, dptr, align)
|
|
return impl_ret_borrowed(context, builder, typ, res)
|
|
|
|
|
|
@lower_setattr_generic(types.Record)
|
|
def record_setattr(context, builder, sig, args, attr):
|
|
"""
|
|
Generic setattr() implementation for records: set the given record member.
|
|
"""
|
|
typ, valty = sig.args
|
|
target, val = args
|
|
|
|
context.sentry_record_alignment(typ, attr)
|
|
offset = typ.offset(attr)
|
|
elemty = typ.typeof(attr)
|
|
|
|
if isinstance(elemty, types.NestedArray):
|
|
# Copy the data from the RHS into the nested array
|
|
val_struct = cgutils.create_struct_proxy(valty)(context, builder,
|
|
value=args[1])
|
|
src = val_struct.data
|
|
dest = cgutils.get_record_member(builder, target, offset,
|
|
src.type.pointee)
|
|
cgutils.memcpy(builder, dest, src,
|
|
context.get_constant(types.intp, elemty.nitems))
|
|
else:
|
|
# Set the given scalar record member
|
|
dptr = cgutils.get_record_member(builder, target, offset,
|
|
context.get_data_type(elemty))
|
|
val = context.cast(builder, val, valty, elemty)
|
|
align = None if typ.aligned else 1
|
|
context.pack_value(builder, elemty, val, dptr, align=align)
|
|
|
|
|
|
@lower_builtin('static_getitem', types.Record, types.StringLiteral)
|
|
def record_static_getitem_str(context, builder, sig, args):
|
|
"""
|
|
Record.__getitem__ redirects to getattr()
|
|
"""
|
|
impl = context.get_getattr(sig.args[0], args[1])
|
|
return impl(context, builder, sig.args[0], args[0], args[1])
|
|
|
|
|
|
@lower_builtin('static_getitem', types.Record, types.IntegerLiteral)
|
|
def record_static_getitem_int(context, builder, sig, args):
|
|
"""
|
|
Record.__getitem__ redirects to getattr()
|
|
"""
|
|
idx = sig.args[1].literal_value
|
|
fields = list(sig.args[0].fields)
|
|
ll_field = context.insert_const_string(builder.module, fields[idx])
|
|
impl = context.get_getattr(sig.args[0], ll_field)
|
|
return impl(context, builder, sig.args[0], args[0], fields[idx])
|
|
|
|
|
|
@lower_builtin('static_setitem', types.Record, types.StringLiteral, types.Any)
|
|
def record_static_setitem_str(context, builder, sig, args):
|
|
"""
|
|
Record.__setitem__ redirects to setattr()
|
|
"""
|
|
recty, _, valty = sig.args
|
|
rec, idx, val = args
|
|
getattr_sig = signature(sig.return_type, recty, valty)
|
|
impl = context.get_setattr(idx, getattr_sig)
|
|
assert impl is not None
|
|
return impl(builder, (rec, val))
|
|
|
|
|
|
@lower_builtin('static_setitem', types.Record, types.IntegerLiteral, types.Any)
|
|
def record_static_setitem_int(context, builder, sig, args):
|
|
"""
|
|
Record.__setitem__ redirects to setattr()
|
|
"""
|
|
recty, _, valty = sig.args
|
|
rec, idx, val = args
|
|
getattr_sig = signature(sig.return_type, recty, valty)
|
|
fields = list(sig.args[0].fields)
|
|
impl = context.get_setattr(fields[idx], getattr_sig)
|
|
assert impl is not None
|
|
return impl(builder, (rec, val))
|
|
|
|
|
|
# ------------------------------------------------------------------------------
|
|
# Constant arrays and records
|
|
|
|
|
|
@lower_constant(types.Array)
|
|
def constant_array(context, builder, ty, pyval):
|
|
"""
|
|
Create a constant array (mechanism is target-dependent).
|
|
"""
|
|
return context.make_constant_array(builder, ty, pyval)
|
|
|
|
|
|
@lower_constant(types.Record)
|
|
def constant_record(context, builder, ty, pyval):
|
|
"""
|
|
Create a record constant as a stack-allocated array of bytes.
|
|
"""
|
|
lty = ir.ArrayType(ir.IntType(8), pyval.nbytes)
|
|
val = lty(bytearray(pyval.tostring()))
|
|
return cgutils.alloca_once_value(builder, val)
|
|
|
|
|
|
@lower_constant(types.Bytes)
|
|
def constant_bytes(context, builder, ty, pyval):
|
|
"""
|
|
Create a constant array from bytes (mechanism is target-dependent).
|
|
"""
|
|
buf = np.array(bytearray(pyval), dtype=np.uint8)
|
|
return context.make_constant_array(builder, ty, buf)
|
|
|
|
# ------------------------------------------------------------------------------
|
|
# Comparisons
|
|
|
|
|
|
@lower_builtin(operator.is_, types.Array, types.Array)
|
|
def array_is(context, builder, sig, args):
|
|
aty, bty = sig.args
|
|
if aty != bty:
|
|
return cgutils.false_bit
|
|
|
|
def array_is_impl(a, b):
|
|
return (a.shape == b.shape and
|
|
a.strides == b.strides and
|
|
a.ctypes.data == b.ctypes.data)
|
|
|
|
return context.compile_internal(builder, array_is_impl, sig, args)
|
|
|
|
# ------------------------------------------------------------------------------
|
|
# Hash
|
|
|
|
|
|
@overload_attribute(types.Array, "__hash__")
|
|
def ol_array_hash(arr):
|
|
return lambda arr: None
|
|
|
|
|
|
# ------------------------------------------------------------------------------
|
|
# builtin `np.flat` implementation
|
|
|
|
def make_array_flat_cls(flatiterty):
|
|
"""
|
|
Return the Structure representation of the given *flatiterty* (an
|
|
instance of types.NumpyFlatType).
|
|
"""
|
|
return _make_flattening_iter_cls(flatiterty, 'flat')
|
|
|
|
|
|
def make_array_ndenumerate_cls(nditerty):
|
|
"""
|
|
Return the Structure representation of the given *nditerty* (an
|
|
instance of types.NumpyNdEnumerateType).
|
|
"""
|
|
return _make_flattening_iter_cls(nditerty, 'ndenumerate')
|
|
|
|
|
|
def _increment_indices(context, builder, ndim, shape, indices, end_flag=None,
|
|
loop_continue=None, loop_break=None):
|
|
zero = context.get_constant(types.intp, 0)
|
|
|
|
bbend = builder.append_basic_block('end_increment')
|
|
|
|
if end_flag is not None:
|
|
builder.store(cgutils.false_byte, end_flag)
|
|
|
|
for dim in reversed(range(ndim)):
|
|
idxptr = cgutils.gep_inbounds(builder, indices, dim)
|
|
idx = cgutils.increment_index(builder, builder.load(idxptr))
|
|
|
|
count = shape[dim]
|
|
in_bounds = builder.icmp_signed('<', idx, count)
|
|
with cgutils.if_likely(builder, in_bounds):
|
|
# New index is still in bounds
|
|
builder.store(idx, idxptr)
|
|
if loop_continue is not None:
|
|
loop_continue(dim)
|
|
builder.branch(bbend)
|
|
# Index out of bounds => reset it and proceed it to outer index
|
|
builder.store(zero, idxptr)
|
|
if loop_break is not None:
|
|
loop_break(dim)
|
|
|
|
if end_flag is not None:
|
|
builder.store(cgutils.true_byte, end_flag)
|
|
builder.branch(bbend)
|
|
|
|
builder.position_at_end(bbend)
|
|
|
|
|
|
def _increment_indices_array(context, builder, arrty, arr, indices,
|
|
end_flag=None):
|
|
shape = cgutils.unpack_tuple(builder, arr.shape, arrty.ndim)
|
|
_increment_indices(context, builder, arrty.ndim, shape, indices, end_flag)
|
|
|
|
|
|
def make_nditer_cls(nditerty):
|
|
"""
|
|
Return the Structure representation of the given *nditerty* (an
|
|
instance of types.NumpyNdIterType).
|
|
"""
|
|
ndim = nditerty.ndim
|
|
layout = nditerty.layout
|
|
narrays = len(nditerty.arrays)
|
|
nshapes = ndim if nditerty.need_shaped_indexing else 1
|
|
|
|
class BaseSubIter(object):
|
|
"""
|
|
Base class for sub-iterators of a nditer() instance.
|
|
"""
|
|
|
|
def __init__(self, nditer, member_name, start_dim, end_dim):
|
|
self.nditer = nditer
|
|
self.member_name = member_name
|
|
self.start_dim = start_dim
|
|
self.end_dim = end_dim
|
|
self.ndim = end_dim - start_dim
|
|
|
|
def set_member_ptr(self, ptr):
|
|
setattr(self.nditer, self.member_name, ptr)
|
|
|
|
@functools.cached_property
|
|
def member_ptr(self):
|
|
return getattr(self.nditer, self.member_name)
|
|
|
|
def init_specific(self, context, builder):
|
|
pass
|
|
|
|
def loop_continue(self, context, builder, logical_dim):
|
|
pass
|
|
|
|
def loop_break(self, context, builder, logical_dim):
|
|
pass
|
|
|
|
class FlatSubIter(BaseSubIter):
|
|
"""
|
|
Sub-iterator walking a contiguous array in physical order, with
|
|
support for broadcasting (the index is reset on the outer dimension).
|
|
"""
|
|
|
|
def init_specific(self, context, builder):
|
|
zero = context.get_constant(types.intp, 0)
|
|
self.set_member_ptr(cgutils.alloca_once_value(builder, zero))
|
|
|
|
def compute_pointer(self, context, builder, indices, arrty, arr):
|
|
index = builder.load(self.member_ptr)
|
|
return builder.gep(arr.data, [index])
|
|
|
|
def loop_continue(self, context, builder, logical_dim):
|
|
if logical_dim == self.ndim - 1:
|
|
# Only increment index inside innermost logical dimension
|
|
index = builder.load(self.member_ptr)
|
|
index = cgutils.increment_index(builder, index)
|
|
builder.store(index, self.member_ptr)
|
|
|
|
def loop_break(self, context, builder, logical_dim):
|
|
if logical_dim == 0:
|
|
# At the exit of outermost logical dimension, reset index
|
|
zero = context.get_constant(types.intp, 0)
|
|
builder.store(zero, self.member_ptr)
|
|
elif logical_dim == self.ndim - 1:
|
|
# Inside innermost logical dimension, increment index
|
|
index = builder.load(self.member_ptr)
|
|
index = cgutils.increment_index(builder, index)
|
|
builder.store(index, self.member_ptr)
|
|
|
|
class TrivialFlatSubIter(BaseSubIter):
|
|
"""
|
|
Sub-iterator walking a contiguous array in physical order,
|
|
*without* support for broadcasting.
|
|
"""
|
|
|
|
def init_specific(self, context, builder):
|
|
assert not nditerty.need_shaped_indexing
|
|
|
|
def compute_pointer(self, context, builder, indices, arrty, arr):
|
|
assert len(indices) <= 1, len(indices)
|
|
return builder.gep(arr.data, indices)
|
|
|
|
class IndexedSubIter(BaseSubIter):
|
|
"""
|
|
Sub-iterator walking an array in logical order.
|
|
"""
|
|
|
|
def compute_pointer(self, context, builder, indices, arrty, arr):
|
|
assert len(indices) == self.ndim
|
|
return cgutils.get_item_pointer(context, builder, arrty, arr,
|
|
indices, wraparound=False)
|
|
|
|
class ZeroDimSubIter(BaseSubIter):
|
|
"""
|
|
Sub-iterator "walking" a 0-d array.
|
|
"""
|
|
|
|
def compute_pointer(self, context, builder, indices, arrty, arr):
|
|
return arr.data
|
|
|
|
class ScalarSubIter(BaseSubIter):
|
|
"""
|
|
Sub-iterator "walking" a scalar value.
|
|
"""
|
|
|
|
def compute_pointer(self, context, builder, indices, arrty, arr):
|
|
return arr
|
|
|
|
class NdIter(cgutils.create_struct_proxy(nditerty)):
|
|
"""
|
|
.nditer() implementation.
|
|
|
|
Note: 'F' layout means the shape is iterated in reverse logical order,
|
|
so indices and shapes arrays have to be reversed as well.
|
|
"""
|
|
|
|
@functools.cached_property
|
|
def subiters(self):
|
|
l = []
|
|
factories = {'flat': FlatSubIter if nditerty.need_shaped_indexing
|
|
else TrivialFlatSubIter,
|
|
'indexed': IndexedSubIter,
|
|
'0d': ZeroDimSubIter,
|
|
'scalar': ScalarSubIter,
|
|
}
|
|
for i, sub in enumerate(nditerty.indexers):
|
|
kind, start_dim, end_dim, _ = sub
|
|
member_name = 'index%d' % i
|
|
factory = factories[kind]
|
|
l.append(factory(self, member_name, start_dim, end_dim))
|
|
return l
|
|
|
|
def init_specific(self, context, builder, arrtys, arrays):
|
|
"""
|
|
Initialize the nditer() instance for the specific array inputs.
|
|
"""
|
|
zero = context.get_constant(types.intp, 0)
|
|
|
|
# Store inputs
|
|
self.arrays = context.make_tuple(builder, types.Tuple(arrtys),
|
|
arrays)
|
|
# Create slots for scalars
|
|
for i, ty in enumerate(arrtys):
|
|
if not isinstance(ty, types.Array):
|
|
member_name = 'scalar%d' % i
|
|
# XXX as_data()?
|
|
slot = cgutils.alloca_once_value(builder, arrays[i])
|
|
setattr(self, member_name, slot)
|
|
|
|
arrays = self._arrays_or_scalars(context, builder, arrtys, arrays)
|
|
|
|
# Extract iterator shape (the shape of the most-dimensional input)
|
|
main_shape_ty = types.UniTuple(types.intp, ndim)
|
|
main_shape = None
|
|
main_nitems = None
|
|
for i, arrty in enumerate(arrtys):
|
|
if isinstance(arrty, types.Array) and arrty.ndim == ndim:
|
|
main_shape = arrays[i].shape
|
|
main_nitems = arrays[i].nitems
|
|
break
|
|
else:
|
|
# Only scalar inputs => synthesize a dummy shape
|
|
assert ndim == 0
|
|
main_shape = context.make_tuple(builder, main_shape_ty, ())
|
|
main_nitems = context.get_constant(types.intp, 1)
|
|
|
|
# Validate shapes of array inputs
|
|
def check_shape(shape, main_shape):
|
|
n = len(shape)
|
|
for i in range(n):
|
|
if shape[i] != main_shape[len(main_shape) - n + i]:
|
|
raise ValueError("nditer(): operands could not be "
|
|
"broadcast together")
|
|
|
|
for arrty, arr in zip(arrtys, arrays):
|
|
if isinstance(arrty, types.Array) and arrty.ndim > 0:
|
|
sig = signature(types.none,
|
|
types.UniTuple(types.intp, arrty.ndim),
|
|
main_shape_ty)
|
|
context.compile_internal(builder, check_shape,
|
|
sig, (arr.shape, main_shape))
|
|
|
|
# Compute shape and size
|
|
shapes = cgutils.unpack_tuple(builder, main_shape)
|
|
if layout == 'F':
|
|
shapes = shapes[::-1]
|
|
|
|
# If shape is empty, mark iterator exhausted
|
|
shape_is_empty = builder.icmp_signed('==', main_nitems, zero)
|
|
exhausted = builder.select(shape_is_empty, cgutils.true_byte,
|
|
cgutils.false_byte)
|
|
|
|
if not nditerty.need_shaped_indexing:
|
|
# Flatten shape to make iteration faster on small innermost
|
|
# dimensions (e.g. a (100000, 3) shape)
|
|
shapes = (main_nitems,)
|
|
assert len(shapes) == nshapes
|
|
|
|
indices = cgutils.alloca_once(builder, zero.type, size=nshapes)
|
|
for dim in range(nshapes):
|
|
idxptr = cgutils.gep_inbounds(builder, indices, dim)
|
|
builder.store(zero, idxptr)
|
|
|
|
self.indices = indices
|
|
self.shape = cgutils.pack_array(builder, shapes, zero.type)
|
|
self.exhausted = cgutils.alloca_once_value(builder, exhausted)
|
|
|
|
# Initialize subiterators
|
|
for subiter in self.subiters:
|
|
subiter.init_specific(context, builder)
|
|
|
|
def iternext_specific(self, context, builder, result):
|
|
"""
|
|
Compute next iteration of the nditer() instance.
|
|
"""
|
|
bbend = builder.append_basic_block('end')
|
|
|
|
# Branch early if exhausted
|
|
exhausted = cgutils.as_bool_bit(builder,
|
|
builder.load(self.exhausted))
|
|
with cgutils.if_unlikely(builder, exhausted):
|
|
result.set_valid(False)
|
|
builder.branch(bbend)
|
|
|
|
arrtys = nditerty.arrays
|
|
arrays = cgutils.unpack_tuple(builder, self.arrays)
|
|
arrays = self._arrays_or_scalars(context, builder, arrtys, arrays)
|
|
indices = self.indices
|
|
|
|
# Compute iterated results
|
|
result.set_valid(True)
|
|
views = self._make_views(context, builder, indices, arrtys, arrays)
|
|
views = [v._getvalue() for v in views]
|
|
if len(views) == 1:
|
|
result.yield_(views[0])
|
|
else:
|
|
result.yield_(context.make_tuple(builder, nditerty.yield_type,
|
|
views))
|
|
|
|
shape = cgutils.unpack_tuple(builder, self.shape)
|
|
_increment_indices(context, builder, len(shape), shape,
|
|
indices, self.exhausted,
|
|
functools.partial(self._loop_continue,
|
|
context,
|
|
builder),
|
|
functools.partial(self._loop_break,
|
|
context,
|
|
builder),
|
|
)
|
|
|
|
builder.branch(bbend)
|
|
builder.position_at_end(bbend)
|
|
|
|
def _loop_continue(self, context, builder, dim):
|
|
for sub in self.subiters:
|
|
if sub.start_dim <= dim < sub.end_dim:
|
|
sub.loop_continue(context, builder, dim - sub.start_dim)
|
|
|
|
def _loop_break(self, context, builder, dim):
|
|
for sub in self.subiters:
|
|
if sub.start_dim <= dim < sub.end_dim:
|
|
sub.loop_break(context, builder, dim - sub.start_dim)
|
|
|
|
def _make_views(self, context, builder, indices, arrtys, arrays):
|
|
"""
|
|
Compute the views to be yielded.
|
|
"""
|
|
views = [None] * narrays
|
|
indexers = nditerty.indexers
|
|
subiters = self.subiters
|
|
rettys = nditerty.yield_type
|
|
if isinstance(rettys, types.BaseTuple):
|
|
rettys = list(rettys)
|
|
else:
|
|
rettys = [rettys]
|
|
indices = [builder.load(cgutils.gep_inbounds(builder, indices, i))
|
|
for i in range(nshapes)]
|
|
|
|
for sub, subiter in zip(indexers, subiters):
|
|
_, _, _, array_indices = sub
|
|
sub_indices = indices[subiter.start_dim:subiter.end_dim]
|
|
if layout == 'F':
|
|
sub_indices = sub_indices[::-1]
|
|
for i in array_indices:
|
|
assert views[i] is None
|
|
views[i] = self._make_view(context, builder, sub_indices,
|
|
rettys[i],
|
|
arrtys[i], arrays[i], subiter)
|
|
assert all(v for v in views)
|
|
return views
|
|
|
|
def _make_view(self, context, builder, indices, retty, arrty, arr,
|
|
subiter):
|
|
"""
|
|
Compute a 0d view for a given input array.
|
|
"""
|
|
assert isinstance(retty, types.Array) and retty.ndim == 0
|
|
|
|
ptr = subiter.compute_pointer(context, builder, indices, arrty, arr)
|
|
view = context.make_array(retty)(context, builder)
|
|
|
|
itemsize = get_itemsize(context, retty)
|
|
shape = context.make_tuple(builder, types.UniTuple(types.intp, 0),
|
|
())
|
|
strides = context.make_tuple(builder, types.UniTuple(types.intp, 0),
|
|
())
|
|
# HACK: meminfo=None avoids expensive refcounting operations
|
|
# on ephemeral views
|
|
populate_array(view, ptr, shape, strides, itemsize, meminfo=None)
|
|
return view
|
|
|
|
def _arrays_or_scalars(self, context, builder, arrtys, arrays):
|
|
# Return a list of either array structures or pointers to
|
|
# scalar slots
|
|
l = []
|
|
for i, (arrty, arr) in enumerate(zip(arrtys, arrays)):
|
|
if isinstance(arrty, types.Array):
|
|
l.append(context.make_array(arrty)(context,
|
|
builder,
|
|
value=arr))
|
|
else:
|
|
l.append(getattr(self, "scalar%d" % i))
|
|
return l
|
|
|
|
return NdIter
|
|
|
|
|
|
def make_ndindex_cls(nditerty):
|
|
"""
|
|
Return the Structure representation of the given *nditerty* (an
|
|
instance of types.NumpyNdIndexType).
|
|
"""
|
|
ndim = nditerty.ndim
|
|
|
|
class NdIndexIter(cgutils.create_struct_proxy(nditerty)):
|
|
"""
|
|
.ndindex() implementation.
|
|
"""
|
|
|
|
def init_specific(self, context, builder, shapes):
|
|
zero = context.get_constant(types.intp, 0)
|
|
indices = cgutils.alloca_once(builder, zero.type,
|
|
size=context.get_constant(types.intp,
|
|
ndim))
|
|
exhausted = cgutils.alloca_once_value(builder, cgutils.false_byte)
|
|
|
|
for dim in range(ndim):
|
|
idxptr = cgutils.gep_inbounds(builder, indices, dim)
|
|
builder.store(zero, idxptr)
|
|
# 0-sized dimensions really indicate an empty array,
|
|
# but we have to catch that condition early to avoid
|
|
# a bug inside the iteration logic.
|
|
dim_size = shapes[dim]
|
|
dim_is_empty = builder.icmp_unsigned('==', dim_size, zero)
|
|
with cgutils.if_unlikely(builder, dim_is_empty):
|
|
builder.store(cgutils.true_byte, exhausted)
|
|
|
|
self.indices = indices
|
|
self.exhausted = exhausted
|
|
self.shape = cgutils.pack_array(builder, shapes, zero.type)
|
|
|
|
def iternext_specific(self, context, builder, result):
|
|
zero = context.get_constant(types.intp, 0)
|
|
|
|
bbend = builder.append_basic_block('end')
|
|
|
|
exhausted = cgutils.as_bool_bit(builder,
|
|
builder.load(self.exhausted))
|
|
with cgutils.if_unlikely(builder, exhausted):
|
|
result.set_valid(False)
|
|
builder.branch(bbend)
|
|
|
|
indices = [builder.load(cgutils.gep_inbounds(builder,
|
|
self.indices,
|
|
dim))
|
|
for dim in range(ndim)]
|
|
for load in indices:
|
|
mark_positive(builder, load)
|
|
|
|
result.yield_(cgutils.pack_array(builder, indices, zero.type))
|
|
result.set_valid(True)
|
|
|
|
shape = cgutils.unpack_tuple(builder, self.shape, ndim)
|
|
_increment_indices(context, builder, ndim, shape,
|
|
self.indices, self.exhausted)
|
|
|
|
builder.branch(bbend)
|
|
builder.position_at_end(bbend)
|
|
|
|
return NdIndexIter
|
|
|
|
|
|
def _make_flattening_iter_cls(flatiterty, kind):
|
|
assert kind in ('flat', 'ndenumerate')
|
|
|
|
array_type = flatiterty.array_type
|
|
|
|
if array_type.layout == 'C':
|
|
class CContiguousFlatIter(cgutils.create_struct_proxy(flatiterty)):
|
|
"""
|
|
.flat() / .ndenumerate() implementation for C-contiguous arrays.
|
|
"""
|
|
|
|
def init_specific(self, context, builder, arrty, arr):
|
|
zero = context.get_constant(types.intp, 0)
|
|
self.index = cgutils.alloca_once_value(builder, zero)
|
|
# We can't trust strides[-1] to always contain the right
|
|
# step value, see
|
|
# http://docs.scipy.org/doc/numpy-dev/release.html#npy-relaxed-strides-checking # noqa: E501
|
|
self.stride = arr.itemsize
|
|
|
|
if kind == 'ndenumerate':
|
|
# Zero-initialize the indices array.
|
|
indices = cgutils.alloca_once(
|
|
builder, zero.type,
|
|
size=context.get_constant(types.intp, arrty.ndim))
|
|
|
|
for dim in range(arrty.ndim):
|
|
idxptr = cgutils.gep_inbounds(builder, indices, dim)
|
|
builder.store(zero, idxptr)
|
|
|
|
self.indices = indices
|
|
|
|
# NOTE: Using gep() instead of explicit pointer addition helps
|
|
# LLVM vectorize the loop (since the stride is known and
|
|
# constant). This is not possible in the non-contiguous case,
|
|
# where the strides are unknown at compile-time.
|
|
|
|
def iternext_specific(self, context, builder, arrty, arr, result):
|
|
ndim = arrty.ndim
|
|
nitems = arr.nitems
|
|
|
|
index = builder.load(self.index)
|
|
is_valid = builder.icmp_signed('<', index, nitems)
|
|
result.set_valid(is_valid)
|
|
|
|
with cgutils.if_likely(builder, is_valid):
|
|
ptr = builder.gep(arr.data, [index])
|
|
value = load_item(context, builder, arrty, ptr)
|
|
if kind == 'flat':
|
|
result.yield_(value)
|
|
else:
|
|
# ndenumerate(): fetch and increment indices
|
|
indices = self.indices
|
|
idxvals = [builder.load(cgutils.gep_inbounds(builder,
|
|
indices,
|
|
dim))
|
|
for dim in range(ndim)]
|
|
idxtuple = cgutils.pack_array(builder, idxvals)
|
|
result.yield_(
|
|
cgutils.make_anonymous_struct(builder,
|
|
[idxtuple, value]))
|
|
_increment_indices_array(context, builder, arrty,
|
|
arr, indices)
|
|
|
|
index = cgutils.increment_index(builder, index)
|
|
builder.store(index, self.index)
|
|
|
|
def getitem(self, context, builder, arrty, arr, index):
|
|
ptr = builder.gep(arr.data, [index])
|
|
return load_item(context, builder, arrty, ptr)
|
|
|
|
def setitem(self, context, builder, arrty, arr, index, value):
|
|
ptr = builder.gep(arr.data, [index])
|
|
store_item(context, builder, arrty, value, ptr)
|
|
|
|
return CContiguousFlatIter
|
|
|
|
else:
|
|
class FlatIter(cgutils.create_struct_proxy(flatiterty)):
|
|
"""
|
|
Generic .flat() / .ndenumerate() implementation for
|
|
non-contiguous arrays.
|
|
It keeps track of pointers along each dimension in order to
|
|
minimize computations.
|
|
"""
|
|
|
|
def init_specific(self, context, builder, arrty, arr):
|
|
zero = context.get_constant(types.intp, 0)
|
|
data = arr.data
|
|
ndim = arrty.ndim
|
|
shapes = cgutils.unpack_tuple(builder, arr.shape, ndim)
|
|
|
|
indices = cgutils.alloca_once(
|
|
builder, zero.type, size=context.get_constant(types.intp,
|
|
arrty.ndim))
|
|
pointers = cgutils.alloca_once(
|
|
builder, data.type, size=context.get_constant(types.intp,
|
|
arrty.ndim))
|
|
exhausted = cgutils.alloca_once_value(builder,
|
|
cgutils.false_byte)
|
|
|
|
# Initialize indices and pointers with their start values.
|
|
for dim in range(ndim):
|
|
idxptr = cgutils.gep_inbounds(builder, indices, dim)
|
|
ptrptr = cgutils.gep_inbounds(builder, pointers, dim)
|
|
builder.store(data, ptrptr)
|
|
builder.store(zero, idxptr)
|
|
# 0-sized dimensions really indicate an empty array,
|
|
# but we have to catch that condition early to avoid
|
|
# a bug inside the iteration logic (see issue #846).
|
|
dim_size = shapes[dim]
|
|
dim_is_empty = builder.icmp_unsigned('==', dim_size, zero)
|
|
with cgutils.if_unlikely(builder, dim_is_empty):
|
|
builder.store(cgutils.true_byte, exhausted)
|
|
|
|
self.indices = indices
|
|
self.pointers = pointers
|
|
self.exhausted = exhausted
|
|
|
|
def iternext_specific(self, context, builder, arrty, arr, result):
|
|
ndim = arrty.ndim
|
|
shapes = cgutils.unpack_tuple(builder, arr.shape, ndim)
|
|
strides = cgutils.unpack_tuple(builder, arr.strides, ndim)
|
|
indices = self.indices
|
|
pointers = self.pointers
|
|
|
|
zero = context.get_constant(types.intp, 0)
|
|
|
|
bbend = builder.append_basic_block('end')
|
|
|
|
# Catch already computed iterator exhaustion
|
|
is_exhausted = cgutils.as_bool_bit(
|
|
builder, builder.load(self.exhausted))
|
|
with cgutils.if_unlikely(builder, is_exhausted):
|
|
result.set_valid(False)
|
|
builder.branch(bbend)
|
|
result.set_valid(True)
|
|
|
|
# Current pointer inside last dimension
|
|
last_ptr = cgutils.gep_inbounds(builder, pointers, ndim - 1)
|
|
ptr = builder.load(last_ptr)
|
|
value = load_item(context, builder, arrty, ptr)
|
|
if kind == 'flat':
|
|
result.yield_(value)
|
|
else:
|
|
# ndenumerate() => yield (indices, value)
|
|
idxvals = [builder.load(cgutils.gep_inbounds(builder,
|
|
indices,
|
|
dim))
|
|
for dim in range(ndim)]
|
|
idxtuple = cgutils.pack_array(builder, idxvals)
|
|
result.yield_(
|
|
cgutils.make_anonymous_struct(builder,
|
|
[idxtuple, value]))
|
|
|
|
# Update indices and pointers by walking from inner
|
|
# dimension to outer.
|
|
for dim in reversed(range(ndim)):
|
|
idxptr = cgutils.gep_inbounds(builder, indices, dim)
|
|
idx = cgutils.increment_index(builder,
|
|
builder.load(idxptr))
|
|
|
|
count = shapes[dim]
|
|
stride = strides[dim]
|
|
in_bounds = builder.icmp_signed('<', idx, count)
|
|
with cgutils.if_likely(builder, in_bounds):
|
|
# Index is valid => pointer can simply be incremented.
|
|
builder.store(idx, idxptr)
|
|
ptrptr = cgutils.gep_inbounds(builder, pointers, dim)
|
|
ptr = builder.load(ptrptr)
|
|
ptr = cgutils.pointer_add(builder, ptr, stride)
|
|
builder.store(ptr, ptrptr)
|
|
# Reset pointers in inner dimensions
|
|
for inner_dim in range(dim + 1, ndim):
|
|
ptrptr = cgutils.gep_inbounds(builder,
|
|
pointers,
|
|
inner_dim)
|
|
builder.store(ptr, ptrptr)
|
|
builder.branch(bbend)
|
|
# Reset index and continue with next dimension
|
|
builder.store(zero, idxptr)
|
|
|
|
# End of array
|
|
builder.store(cgutils.true_byte, self.exhausted)
|
|
builder.branch(bbend)
|
|
|
|
builder.position_at_end(bbend)
|
|
|
|
def _ptr_for_index(self, context, builder, arrty, arr, index):
|
|
ndim = arrty.ndim
|
|
shapes = cgutils.unpack_tuple(builder, arr.shape, count=ndim)
|
|
strides = cgutils.unpack_tuple(builder, arr.strides, count=ndim)
|
|
|
|
# First convert the flattened index into a regular n-dim index
|
|
indices = []
|
|
for dim in reversed(range(ndim)):
|
|
indices.append(builder.urem(index, shapes[dim]))
|
|
index = builder.udiv(index, shapes[dim])
|
|
indices.reverse()
|
|
|
|
ptr = cgutils.get_item_pointer2(context, builder, arr.data,
|
|
shapes, strides, arrty.layout,
|
|
indices)
|
|
return ptr
|
|
|
|
def getitem(self, context, builder, arrty, arr, index):
|
|
ptr = self._ptr_for_index(context, builder, arrty, arr, index)
|
|
return load_item(context, builder, arrty, ptr)
|
|
|
|
def setitem(self, context, builder, arrty, arr, index, value):
|
|
ptr = self._ptr_for_index(context, builder, arrty, arr, index)
|
|
store_item(context, builder, arrty, value, ptr)
|
|
|
|
return FlatIter
|
|
|
|
|
|
@lower_getattr(types.Array, "flat")
|
|
def make_array_flatiter(context, builder, arrty, arr):
|
|
flatitercls = make_array_flat_cls(types.NumpyFlatType(arrty))
|
|
flatiter = flatitercls(context, builder)
|
|
|
|
flatiter.array = arr
|
|
|
|
arrcls = context.make_array(arrty)
|
|
arr = arrcls(context, builder, ref=flatiter._get_ptr_by_name('array'))
|
|
|
|
flatiter.init_specific(context, builder, arrty, arr)
|
|
|
|
res = flatiter._getvalue()
|
|
return impl_ret_borrowed(context, builder, types.NumpyFlatType(arrty), res)
|
|
|
|
|
|
@lower_builtin('iternext', types.NumpyFlatType)
|
|
@iternext_impl(RefType.BORROWED)
|
|
def iternext_numpy_flatiter(context, builder, sig, args, result):
|
|
[flatiterty] = sig.args
|
|
[flatiter] = args
|
|
|
|
flatitercls = make_array_flat_cls(flatiterty)
|
|
flatiter = flatitercls(context, builder, value=flatiter)
|
|
|
|
arrty = flatiterty.array_type
|
|
arrcls = context.make_array(arrty)
|
|
arr = arrcls(context, builder, value=flatiter.array)
|
|
|
|
flatiter.iternext_specific(context, builder, arrty, arr, result)
|
|
|
|
|
|
@lower_builtin(operator.getitem, types.NumpyFlatType, types.Integer)
|
|
def iternext_numpy_getitem(context, builder, sig, args):
|
|
flatiterty = sig.args[0]
|
|
flatiter, index = args
|
|
|
|
flatitercls = make_array_flat_cls(flatiterty)
|
|
flatiter = flatitercls(context, builder, value=flatiter)
|
|
|
|
arrty = flatiterty.array_type
|
|
arrcls = context.make_array(arrty)
|
|
arr = arrcls(context, builder, value=flatiter.array)
|
|
|
|
res = flatiter.getitem(context, builder, arrty, arr, index)
|
|
return impl_ret_borrowed(context, builder, sig.return_type, res)
|
|
|
|
|
|
@lower_builtin(operator.setitem, types.NumpyFlatType, types.Integer,
|
|
types.Any)
|
|
def iternext_numpy_getitem_any(context, builder, sig, args):
|
|
flatiterty = sig.args[0]
|
|
flatiter, index, value = args
|
|
|
|
flatitercls = make_array_flat_cls(flatiterty)
|
|
flatiter = flatitercls(context, builder, value=flatiter)
|
|
|
|
arrty = flatiterty.array_type
|
|
arrcls = context.make_array(arrty)
|
|
arr = arrcls(context, builder, value=flatiter.array)
|
|
|
|
flatiter.setitem(context, builder, arrty, arr, index, value)
|
|
return context.get_dummy_value()
|
|
|
|
|
|
@lower_builtin(len, types.NumpyFlatType)
|
|
def iternext_numpy_getitem_flat(context, builder, sig, args):
|
|
flatiterty = sig.args[0]
|
|
flatitercls = make_array_flat_cls(flatiterty)
|
|
flatiter = flatitercls(context, builder, value=args[0])
|
|
|
|
arrcls = context.make_array(flatiterty.array_type)
|
|
arr = arrcls(context, builder, value=flatiter.array)
|
|
return arr.nitems
|
|
|
|
|
|
@lower_builtin(np.ndenumerate, types.Array)
|
|
def make_array_ndenumerate(context, builder, sig, args):
|
|
arrty, = sig.args
|
|
arr, = args
|
|
nditercls = make_array_ndenumerate_cls(types.NumpyNdEnumerateType(arrty))
|
|
nditer = nditercls(context, builder)
|
|
|
|
nditer.array = arr
|
|
|
|
arrcls = context.make_array(arrty)
|
|
arr = arrcls(context, builder, ref=nditer._get_ptr_by_name('array'))
|
|
|
|
nditer.init_specific(context, builder, arrty, arr)
|
|
|
|
res = nditer._getvalue()
|
|
return impl_ret_borrowed(context, builder, sig.return_type, res)
|
|
|
|
|
|
@lower_builtin('iternext', types.NumpyNdEnumerateType)
|
|
@iternext_impl(RefType.BORROWED)
|
|
def iternext_numpy_nditer(context, builder, sig, args, result):
|
|
[nditerty] = sig.args
|
|
[nditer] = args
|
|
|
|
nditercls = make_array_ndenumerate_cls(nditerty)
|
|
nditer = nditercls(context, builder, value=nditer)
|
|
|
|
arrty = nditerty.array_type
|
|
arrcls = context.make_array(arrty)
|
|
arr = arrcls(context, builder, value=nditer.array)
|
|
|
|
nditer.iternext_specific(context, builder, arrty, arr, result)
|
|
|
|
|
|
@lower_builtin(pndindex, types.VarArg(types.Integer))
|
|
@lower_builtin(np.ndindex, types.VarArg(types.Integer))
|
|
def make_array_ndindex(context, builder, sig, args):
|
|
"""ndindex(*shape)"""
|
|
shape = [context.cast(builder, arg, argty, types.intp)
|
|
for argty, arg in zip(sig.args, args)]
|
|
|
|
nditercls = make_ndindex_cls(types.NumpyNdIndexType(len(shape)))
|
|
nditer = nditercls(context, builder)
|
|
nditer.init_specific(context, builder, shape)
|
|
|
|
res = nditer._getvalue()
|
|
return impl_ret_borrowed(context, builder, sig.return_type, res)
|
|
|
|
|
|
@lower_builtin(pndindex, types.BaseTuple)
|
|
@lower_builtin(np.ndindex, types.BaseTuple)
|
|
def make_array_ndindex_tuple(context, builder, sig, args):
|
|
"""ndindex(shape)"""
|
|
ndim = sig.return_type.ndim
|
|
if ndim > 0:
|
|
idxty = sig.args[0].dtype
|
|
tup = args[0]
|
|
|
|
shape = cgutils.unpack_tuple(builder, tup, ndim)
|
|
shape = [context.cast(builder, idx, idxty, types.intp)
|
|
for idx in shape]
|
|
else:
|
|
shape = []
|
|
|
|
nditercls = make_ndindex_cls(types.NumpyNdIndexType(len(shape)))
|
|
nditer = nditercls(context, builder)
|
|
nditer.init_specific(context, builder, shape)
|
|
|
|
res = nditer._getvalue()
|
|
return impl_ret_borrowed(context, builder, sig.return_type, res)
|
|
|
|
|
|
@lower_builtin('iternext', types.NumpyNdIndexType)
|
|
@iternext_impl(RefType.BORROWED)
|
|
def iternext_numpy_ndindex(context, builder, sig, args, result):
|
|
[nditerty] = sig.args
|
|
[nditer] = args
|
|
|
|
nditercls = make_ndindex_cls(nditerty)
|
|
nditer = nditercls(context, builder, value=nditer)
|
|
|
|
nditer.iternext_specific(context, builder, result)
|
|
|
|
|
|
@lower_builtin(np.nditer, types.Any)
|
|
def make_array_nditer(context, builder, sig, args):
|
|
"""
|
|
nditer(...)
|
|
"""
|
|
nditerty = sig.return_type
|
|
arrtys = nditerty.arrays
|
|
|
|
if isinstance(sig.args[0], types.BaseTuple):
|
|
arrays = cgutils.unpack_tuple(builder, args[0])
|
|
else:
|
|
arrays = [args[0]]
|
|
|
|
nditer = make_nditer_cls(nditerty)(context, builder)
|
|
nditer.init_specific(context, builder, arrtys, arrays)
|
|
|
|
res = nditer._getvalue()
|
|
return impl_ret_borrowed(context, builder, nditerty, res)
|
|
|
|
|
|
@lower_builtin('iternext', types.NumpyNdIterType)
|
|
@iternext_impl(RefType.BORROWED)
|
|
def iternext_numpy_nditer2(context, builder, sig, args, result):
|
|
[nditerty] = sig.args
|
|
[nditer] = args
|
|
|
|
nditer = make_nditer_cls(nditerty)(context, builder, value=nditer)
|
|
nditer.iternext_specific(context, builder, result)
|
|
|
|
|
|
@lower_builtin(operator.eq, types.DType, types.DType)
|
|
def dtype_eq_impl(context, builder, sig, args):
|
|
arg1, arg2 = sig.args
|
|
res = ir.Constant(ir.IntType(1), int(arg1 == arg2))
|
|
return impl_ret_untracked(context, builder, sig.return_type, res)
|
|
|
|
|
|
# ------------------------------------------------------------------------------
|
|
# Numpy array constructors
|
|
|
|
def _empty_nd_impl(context, builder, arrtype, shapes):
|
|
"""Utility function used for allocating a new array during LLVM code
|
|
generation (lowering). Given a target context, builder, array
|
|
type, and a tuple or list of lowered dimension sizes, returns a
|
|
LLVM value pointing at a Numba runtime allocated array.
|
|
"""
|
|
arycls = make_array(arrtype)
|
|
ary = arycls(context, builder)
|
|
|
|
datatype = context.get_data_type(arrtype.dtype)
|
|
itemsize = context.get_constant(types.intp, get_itemsize(context, arrtype))
|
|
|
|
# compute array length
|
|
arrlen = context.get_constant(types.intp, 1)
|
|
overflow = Constant(ir.IntType(1), 0)
|
|
for s in shapes:
|
|
arrlen_mult = builder.smul_with_overflow(arrlen, s)
|
|
arrlen = builder.extract_value(arrlen_mult, 0)
|
|
overflow = builder.or_(
|
|
overflow, builder.extract_value(arrlen_mult, 1)
|
|
)
|
|
|
|
if arrtype.ndim == 0:
|
|
strides = ()
|
|
elif arrtype.layout == 'C':
|
|
strides = [itemsize]
|
|
for dimension_size in reversed(shapes[1:]):
|
|
strides.append(builder.mul(strides[-1], dimension_size))
|
|
strides = tuple(reversed(strides))
|
|
elif arrtype.layout == 'F':
|
|
strides = [itemsize]
|
|
for dimension_size in shapes[:-1]:
|
|
strides.append(builder.mul(strides[-1], dimension_size))
|
|
strides = tuple(strides)
|
|
else:
|
|
raise NotImplementedError(
|
|
"Don't know how to allocate array with layout '{0}'.".format(
|
|
arrtype.layout))
|
|
|
|
# Check overflow, numpy also does this after checking order
|
|
allocsize_mult = builder.smul_with_overflow(arrlen, itemsize)
|
|
allocsize = builder.extract_value(allocsize_mult, 0)
|
|
overflow = builder.or_(overflow, builder.extract_value(allocsize_mult, 1))
|
|
|
|
with builder.if_then(overflow, likely=False):
|
|
# Raise same error as numpy, see:
|
|
# https://github.com/numpy/numpy/blob/2a488fe76a0f732dc418d03b452caace161673da/numpy/core/src/multiarray/ctors.c#L1095-L1101 # noqa: E501
|
|
context.call_conv.return_user_exc(
|
|
builder, ValueError,
|
|
("array is too big; `arr.size * arr.dtype.itemsize` is larger than"
|
|
" the maximum possible size.",)
|
|
)
|
|
|
|
dtype = arrtype.dtype
|
|
align_val = context.get_preferred_array_alignment(dtype)
|
|
align = context.get_constant(types.uint32, align_val)
|
|
args = (context.get_dummy_value(), allocsize, align)
|
|
|
|
mip = types.MemInfoPointer(types.voidptr)
|
|
arytypeclass = types.TypeRef(type(arrtype))
|
|
argtypes = signature(mip, arytypeclass, types.intp, types.uint32)
|
|
|
|
meminfo = context.compile_internal(builder, _call_allocator, argtypes, args)
|
|
data = context.nrt.meminfo_data(builder, meminfo)
|
|
|
|
intp_t = context.get_value_type(types.intp)
|
|
shape_array = cgutils.pack_array(builder, shapes, ty=intp_t)
|
|
strides_array = cgutils.pack_array(builder, strides, ty=intp_t)
|
|
|
|
populate_array(ary,
|
|
data=builder.bitcast(data, datatype.as_pointer()),
|
|
shape=shape_array,
|
|
strides=strides_array,
|
|
itemsize=itemsize,
|
|
meminfo=meminfo)
|
|
|
|
return ary
|
|
|
|
|
|
@overload_classmethod(types.Array, "_allocate")
|
|
def _ol_array_allocate(cls, allocsize, align):
|
|
"""Implements a Numba-only default target (cpu) classmethod on the array
|
|
type.
|
|
"""
|
|
def impl(cls, allocsize, align):
|
|
return intrin_alloc(allocsize, align)
|
|
return impl
|
|
|
|
|
|
def _call_allocator(arrtype, size, align):
|
|
"""Trampoline to call the intrinsic used for allocation
|
|
"""
|
|
return arrtype._allocate(size, align)
|
|
|
|
|
|
@intrinsic
|
|
def intrin_alloc(typingctx, allocsize, align):
|
|
"""Intrinsic to call into the allocator for Array
|
|
"""
|
|
def codegen(context, builder, signature, args):
|
|
[allocsize, align] = args
|
|
meminfo = context.nrt.meminfo_alloc_aligned(builder, allocsize, align)
|
|
return meminfo
|
|
|
|
mip = types.MemInfoPointer(types.voidptr) # return untyped pointer
|
|
sig = signature(mip, allocsize, align)
|
|
return sig, codegen
|
|
|
|
|
|
def _parse_shape(context, builder, ty, val):
|
|
"""
|
|
Parse the shape argument to an array constructor.
|
|
"""
|
|
def safecast_intp(context, builder, src_t, src):
|
|
"""Cast src to intp only if value can be maintained"""
|
|
intp_t = context.get_value_type(types.intp)
|
|
intp_width = intp_t.width
|
|
intp_ir = ir.IntType(intp_width)
|
|
maxval = Constant(intp_ir, ((1 << intp_width - 1) - 1))
|
|
if src_t.width < intp_width:
|
|
res = builder.sext(src, intp_ir)
|
|
elif src_t.width >= intp_width:
|
|
is_larger = builder.icmp_signed(">", src, maxval)
|
|
with builder.if_then(is_larger, likely=False):
|
|
context.call_conv.return_user_exc(
|
|
builder, ValueError,
|
|
("Cannot safely convert value to intp",)
|
|
)
|
|
if src_t.width > intp_width:
|
|
res = builder.trunc(src, intp_ir)
|
|
else:
|
|
res = src
|
|
return res
|
|
|
|
if isinstance(ty, types.Integer):
|
|
ndim = 1
|
|
passed_shapes = [context.cast(builder, val, ty, types.intp)]
|
|
else:
|
|
assert isinstance(ty, types.BaseTuple)
|
|
ndim = ty.count
|
|
passed_shapes = cgutils.unpack_tuple(builder, val, count=ndim)
|
|
|
|
shapes = []
|
|
for s in passed_shapes:
|
|
shapes.append(safecast_intp(context, builder, s.type, s))
|
|
|
|
zero = context.get_constant_generic(builder, types.intp, 0)
|
|
for dim in range(ndim):
|
|
is_neg = builder.icmp_signed('<', shapes[dim], zero)
|
|
with cgutils.if_unlikely(builder, is_neg):
|
|
context.call_conv.return_user_exc(
|
|
builder, ValueError, ("negative dimensions not allowed",)
|
|
)
|
|
|
|
return shapes
|
|
|
|
|
|
def _parse_empty_args(context, builder, sig, args):
|
|
"""
|
|
Parse the arguments of a np.empty(), np.zeros() or np.ones() call.
|
|
"""
|
|
arrshapetype = sig.args[0]
|
|
arrshape = args[0]
|
|
arrtype = sig.return_type
|
|
return arrtype, _parse_shape(context, builder, arrshapetype, arrshape)
|
|
|
|
|
|
def _parse_empty_like_args(context, builder, sig, args):
|
|
"""
|
|
Parse the arguments of a np.empty_like(), np.zeros_like() or
|
|
np.ones_like() call.
|
|
"""
|
|
arytype = sig.args[0]
|
|
if isinstance(arytype, types.Array):
|
|
ary = make_array(arytype)(context, builder, value=args[0])
|
|
shapes = cgutils.unpack_tuple(builder, ary.shape, count=arytype.ndim)
|
|
return sig.return_type, shapes
|
|
else:
|
|
return sig.return_type, ()
|
|
|
|
|
|
def _check_const_str_dtype(fname, dtype):
|
|
if isinstance(dtype, types.UnicodeType):
|
|
msg = f"If np.{fname} dtype is a string it must be a string constant."
|
|
raise errors.TypingError(msg)
|
|
|
|
|
|
@intrinsic
|
|
def numpy_empty_nd(tyctx, ty_shape, ty_dtype, ty_retty_ref):
|
|
ty_retty = ty_retty_ref.instance_type
|
|
sig = ty_retty(ty_shape, ty_dtype, ty_retty_ref)
|
|
|
|
def codegen(cgctx, builder, sig, llargs):
|
|
arrtype, shapes = _parse_empty_args(cgctx, builder, sig, llargs)
|
|
ary = _empty_nd_impl(cgctx, builder, arrtype, shapes)
|
|
return ary._getvalue()
|
|
return sig, codegen
|
|
|
|
|
|
@overload(np.empty)
|
|
def ol_np_empty(shape, dtype=float):
|
|
_check_const_str_dtype("empty", dtype)
|
|
if (dtype is float or
|
|
(isinstance(dtype, types.Function) and dtype.typing_key is float) or
|
|
is_nonelike(dtype)): #default
|
|
nb_dtype = types.double
|
|
else:
|
|
nb_dtype = ty_parse_dtype(dtype)
|
|
|
|
ndim = ty_parse_shape(shape)
|
|
if nb_dtype is not None and ndim is not None:
|
|
retty = types.Array(dtype=nb_dtype, ndim=ndim, layout='C')
|
|
|
|
def impl(shape, dtype=float):
|
|
return numpy_empty_nd(shape, dtype, retty)
|
|
return impl
|
|
else:
|
|
msg = f"Cannot parse input types to function np.empty({shape}, {dtype})"
|
|
raise errors.TypingError(msg)
|
|
|
|
|
|
@intrinsic
|
|
def numpy_empty_like_nd(tyctx, ty_prototype, ty_dtype, ty_retty_ref):
|
|
ty_retty = ty_retty_ref.instance_type
|
|
sig = ty_retty(ty_prototype, ty_dtype, ty_retty_ref)
|
|
|
|
def codegen(cgctx, builder, sig, llargs):
|
|
arrtype, shapes = _parse_empty_like_args(cgctx, builder, sig, llargs)
|
|
ary = _empty_nd_impl(cgctx, builder, arrtype, shapes)
|
|
return ary._getvalue()
|
|
return sig, codegen
|
|
|
|
|
|
@overload(np.empty_like)
|
|
def ol_np_empty_like(arr, dtype=None):
|
|
_check_const_str_dtype("empty_like", dtype)
|
|
if not is_nonelike(dtype):
|
|
nb_dtype = ty_parse_dtype(dtype)
|
|
elif isinstance(arr, types.Array):
|
|
nb_dtype = arr.dtype
|
|
else:
|
|
nb_dtype = arr
|
|
if nb_dtype is not None:
|
|
if isinstance(arr, types.Array):
|
|
layout = arr.layout if arr.layout != 'A' else 'C'
|
|
retty = arr.copy(dtype=nb_dtype, layout=layout, readonly=False)
|
|
else:
|
|
retty = types.Array(nb_dtype, 0, 'C')
|
|
else:
|
|
msg = ("Cannot parse input types to function "
|
|
f"np.empty_like({arr}, {dtype})")
|
|
raise errors.TypingError(msg)
|
|
|
|
def impl(arr, dtype=None):
|
|
return numpy_empty_like_nd(arr, dtype, retty)
|
|
return impl
|
|
|
|
|
|
@intrinsic
|
|
def _zero_fill_array_method(tyctx, self):
|
|
sig = types.none(self)
|
|
|
|
def codegen(cgctx, builder, sig, llargs):
|
|
ary = make_array(sig.args[0])(cgctx, builder, llargs[0])
|
|
cgutils.memset(builder, ary.data, builder.mul(ary.itemsize, ary.nitems),
|
|
0)
|
|
return sig, codegen
|
|
|
|
|
|
@overload_method(types.Array, '_zero_fill')
|
|
def ol_array_zero_fill(self):
|
|
"""Adds a `._zero_fill` method to zero fill an array using memset."""
|
|
def impl(self):
|
|
_zero_fill_array_method(self)
|
|
return impl
|
|
|
|
|
|
@overload(np.zeros)
|
|
def ol_np_zeros(shape, dtype=float):
|
|
_check_const_str_dtype("zeros", dtype)
|
|
|
|
def impl(shape, dtype=float):
|
|
arr = np.empty(shape, dtype=dtype)
|
|
arr._zero_fill()
|
|
return arr
|
|
return impl
|
|
|
|
|
|
@overload(np.zeros_like)
|
|
def ol_np_zeros_like(a, dtype=None):
|
|
_check_const_str_dtype("zeros_like", dtype)
|
|
|
|
# NumPy uses 'a' as the arg name for the array-like
|
|
def impl(a, dtype=None):
|
|
arr = np.empty_like(a, dtype=dtype)
|
|
arr._zero_fill()
|
|
return arr
|
|
return impl
|
|
|
|
|
|
@overload(np.ones_like)
|
|
def ol_np_ones_like(a, dtype=None):
|
|
_check_const_str_dtype("ones_like", dtype)
|
|
|
|
# NumPy uses 'a' as the arg name for the array-like
|
|
def impl(a, dtype=None):
|
|
arr = np.empty_like(a, dtype=dtype)
|
|
arr_flat = arr.flat
|
|
for idx in range(len(arr_flat)):
|
|
arr_flat[idx] = 1
|
|
return arr
|
|
return impl
|
|
|
|
|
|
@overload(np.full)
|
|
def impl_np_full(shape, fill_value, dtype=None):
|
|
_check_const_str_dtype("full", dtype)
|
|
if not is_nonelike(dtype):
|
|
nb_dtype = ty_parse_dtype(dtype)
|
|
else:
|
|
nb_dtype = fill_value
|
|
|
|
def full(shape, fill_value, dtype=None):
|
|
arr = np.empty(shape, nb_dtype)
|
|
arr_flat = arr.flat
|
|
for idx in range(len(arr_flat)):
|
|
arr_flat[idx] = fill_value
|
|
return arr
|
|
return full
|
|
|
|
|
|
@overload(np.full_like)
|
|
def impl_np_full_like(a, fill_value, dtype=None):
|
|
_check_const_str_dtype("full_like", dtype)
|
|
|
|
def full_like(a, fill_value, dtype=None):
|
|
arr = np.empty_like(a, dtype)
|
|
arr_flat = arr.flat
|
|
for idx in range(len(arr_flat)):
|
|
arr_flat[idx] = fill_value
|
|
return arr
|
|
|
|
return full_like
|
|
|
|
|
|
@overload(np.ones)
|
|
def ol_np_ones(shape, dtype=None):
|
|
# for some reason the NumPy default for dtype is None in the source but
|
|
# ends up as np.float64 by definition.
|
|
_check_const_str_dtype("ones", dtype)
|
|
|
|
def impl(shape, dtype=None):
|
|
arr = np.empty(shape, dtype=dtype)
|
|
arr_flat = arr.flat
|
|
for idx in range(len(arr_flat)):
|
|
arr_flat[idx] = 1
|
|
return arr
|
|
return impl
|
|
|
|
|
|
@overload(np.identity)
|
|
def impl_np_identity(n, dtype=None):
|
|
_check_const_str_dtype("identity", dtype)
|
|
if not is_nonelike(dtype):
|
|
nb_dtype = ty_parse_dtype(dtype)
|
|
else:
|
|
nb_dtype = types.double
|
|
|
|
def identity(n, dtype=None):
|
|
arr = np.zeros((n, n), nb_dtype)
|
|
for i in range(n):
|
|
arr[i, i] = 1
|
|
return arr
|
|
return identity
|
|
|
|
|
|
def _eye_none_handler(N, M):
|
|
pass
|
|
|
|
|
|
@extending.overload(_eye_none_handler)
|
|
def _eye_none_handler_impl(N, M):
|
|
if isinstance(M, types.NoneType):
|
|
def impl(N, M):
|
|
return N
|
|
else:
|
|
def impl(N, M):
|
|
return M
|
|
return impl
|
|
|
|
|
|
@extending.overload(np.eye)
|
|
def numpy_eye(N, M=None, k=0, dtype=float):
|
|
|
|
if dtype is None or isinstance(dtype, types.NoneType):
|
|
dt = np.dtype(float)
|
|
elif isinstance(dtype, (types.DTypeSpec, types.Number)):
|
|
# dtype or instance of dtype
|
|
dt = as_dtype(getattr(dtype, 'dtype', dtype))
|
|
else:
|
|
dt = np.dtype(dtype)
|
|
|
|
def impl(N, M=None, k=0, dtype=float):
|
|
_M = _eye_none_handler(N, M)
|
|
arr = np.zeros((N, _M), dt)
|
|
if k >= 0:
|
|
d = min(N, _M - k)
|
|
for i in range(d):
|
|
arr[i, i + k] = 1
|
|
else:
|
|
d = min(N + k, _M)
|
|
for i in range(d):
|
|
arr[i - k, i] = 1
|
|
return arr
|
|
return impl
|
|
|
|
|
|
@overload(np.diag)
|
|
def impl_np_diag(v, k=0):
|
|
if not type_can_asarray(v):
|
|
raise errors.TypingError('The argument "v" must be array-like')
|
|
|
|
if isinstance(v, types.Array):
|
|
if v.ndim not in (1, 2):
|
|
raise errors.NumbaTypeError("Input must be 1- or 2-d.")
|
|
|
|
def diag_impl(v, k=0):
|
|
if v.ndim == 1:
|
|
s = v.shape
|
|
n = s[0] + abs(k)
|
|
ret = np.zeros((n, n), v.dtype)
|
|
if k >= 0:
|
|
for i in range(n - k):
|
|
ret[i, k + i] = v[i]
|
|
else:
|
|
for i in range(n + k):
|
|
ret[i - k, i] = v[i]
|
|
return ret
|
|
else: # 2-d
|
|
rows, cols = v.shape
|
|
if k < 0:
|
|
rows = rows + k
|
|
if k > 0:
|
|
cols = cols - k
|
|
n = max(min(rows, cols), 0)
|
|
ret = np.empty(n, v.dtype)
|
|
if k >= 0:
|
|
for i in range(n):
|
|
ret[i] = v[i, k + i]
|
|
else:
|
|
for i in range(n):
|
|
ret[i] = v[i - k, i]
|
|
return ret
|
|
return diag_impl
|
|
|
|
|
|
@overload(np.indices)
|
|
def numpy_indices(dimensions):
|
|
if not isinstance(dimensions, types.UniTuple):
|
|
msg = 'The argument "dimensions" must be a tuple of integers'
|
|
raise errors.TypingError(msg)
|
|
|
|
if not isinstance(dimensions.dtype, types.Integer):
|
|
msg = 'The argument "dimensions" must be a tuple of integers'
|
|
raise errors.TypingError(msg)
|
|
|
|
N = len(dimensions)
|
|
shape = (1,) * N
|
|
|
|
def impl(dimensions):
|
|
res = np.empty((N,) + dimensions, dtype=np.int64)
|
|
i = 0
|
|
for dim in dimensions:
|
|
idx = np.arange(dim, dtype=np.int64).reshape(
|
|
tuple_setitem(shape, i, dim)
|
|
)
|
|
res[i] = idx
|
|
i += 1
|
|
|
|
return res
|
|
|
|
return impl
|
|
|
|
|
|
@overload(np.diagflat)
|
|
def numpy_diagflat(v, k=0):
|
|
if not type_can_asarray(v):
|
|
msg = 'The argument "v" must be array-like'
|
|
raise errors.TypingError(msg)
|
|
|
|
if not isinstance(k, (int, types.Integer)):
|
|
msg = 'The argument "k" must be an integer'
|
|
raise errors.TypingError(msg)
|
|
|
|
def impl(v, k=0):
|
|
v = np.asarray(v)
|
|
v = v.ravel()
|
|
s = len(v)
|
|
abs_k = abs(k)
|
|
n = s + abs_k
|
|
res = np.zeros((n, n), v.dtype)
|
|
i = np.maximum(0, -k)
|
|
j = np.maximum(0, k)
|
|
for t in range(s):
|
|
res[i + t, j + t] = v[t]
|
|
|
|
return res
|
|
|
|
return impl
|
|
|
|
|
|
@overload(np.take)
|
|
@overload_method(types.Array, 'take')
|
|
def numpy_take(a, indices):
|
|
|
|
if isinstance(a, types.Array) and isinstance(indices, types.Integer):
|
|
def take_impl(a, indices):
|
|
if indices > (a.size - 1) or indices < -a.size:
|
|
raise IndexError("Index out of bounds")
|
|
return a.ravel()[indices]
|
|
return take_impl
|
|
|
|
if all(isinstance(arg, types.Array) for arg in [a, indices]):
|
|
F_order = indices.layout == 'F'
|
|
|
|
def take_impl(a, indices):
|
|
ret = np.empty(indices.size, dtype=a.dtype)
|
|
if F_order:
|
|
walker = indices.copy() # get C order
|
|
else:
|
|
walker = indices
|
|
it = np.nditer(walker)
|
|
i = 0
|
|
flat = a.ravel()
|
|
for x in it:
|
|
if x > (a.size - 1) or x < -a.size:
|
|
raise IndexError("Index out of bounds")
|
|
ret[i] = flat[x]
|
|
i = i + 1
|
|
return ret.reshape(indices.shape)
|
|
return take_impl
|
|
|
|
if isinstance(a, types.Array) and \
|
|
isinstance(indices, (types.List, types.BaseTuple)):
|
|
def take_impl(a, indices):
|
|
convert = np.array(indices)
|
|
ret = np.empty(convert.size, dtype=a.dtype)
|
|
it = np.nditer(convert)
|
|
i = 0
|
|
flat = a.ravel()
|
|
for x in it:
|
|
if x > (a.size - 1) or x < -a.size:
|
|
raise IndexError("Index out of bounds")
|
|
ret[i] = flat[x]
|
|
i = i + 1
|
|
return ret.reshape(convert.shape)
|
|
return take_impl
|
|
|
|
|
|
def _arange_dtype(*args):
|
|
bounds = [a for a in args if not isinstance(a, types.NoneType)]
|
|
|
|
if any(isinstance(a, types.Complex) for a in bounds):
|
|
dtype = types.complex128
|
|
elif any(isinstance(a, types.Float) for a in bounds):
|
|
dtype = types.float64
|
|
else:
|
|
# `np.arange(10).dtype` is always `np.dtype(int)`, aka `np.int_`, which
|
|
# in all released versions of numpy corresponds to the C `long` type.
|
|
# Windows 64 is broken by default here because Numba (as of 0.47) does
|
|
# not differentiate between Python and NumPy integers, so a `typeof(1)`
|
|
# on w64 is `int64`, i.e. `intp`. This means an arange(<some int>) will
|
|
# be typed as arange(int64) and the following will yield int64 opposed
|
|
# to int32. Example: without a load of analysis to work out of the args
|
|
# were wrapped in NumPy int*() calls it's not possible to detect the
|
|
# difference between `np.arange(10)` and `np.arange(np.int64(10)`.
|
|
NPY_TY = getattr(types, "int%s" % (8 * np.dtype(int).itemsize))
|
|
|
|
# unliteral these types such that `max` works.
|
|
unliteral_bounds = [types.unliteral(x) for x in bounds]
|
|
dtype = max(unliteral_bounds + [NPY_TY,])
|
|
|
|
return dtype
|
|
|
|
|
|
@overload(np.arange)
|
|
def np_arange(start, stop=None, step=None, dtype=None):
|
|
if isinstance(stop, types.Optional):
|
|
stop = stop.type
|
|
if isinstance(step, types.Optional):
|
|
step = step.type
|
|
if isinstance(dtype, types.Optional):
|
|
dtype = dtype.type
|
|
|
|
if stop is None:
|
|
stop = types.none
|
|
if step is None:
|
|
step = types.none
|
|
if dtype is None:
|
|
dtype = types.none
|
|
|
|
if (not isinstance(start, types.Number) or
|
|
not isinstance(stop, (types.NoneType, types.Number)) or
|
|
not isinstance(step, (types.NoneType, types.Number)) or
|
|
not isinstance(dtype, (types.NoneType, types.DTypeSpec))):
|
|
|
|
return
|
|
|
|
if isinstance(dtype, types.NoneType):
|
|
true_dtype = _arange_dtype(start, stop, step)
|
|
else:
|
|
true_dtype = dtype.dtype
|
|
|
|
use_complex = any([isinstance(x, types.Complex)
|
|
for x in (start, stop, step)])
|
|
|
|
start_value = getattr(start, "literal_value", None)
|
|
stop_value = getattr(stop, "literal_value", None)
|
|
step_value = getattr(step, "literal_value", None)
|
|
|
|
def impl(start, stop=None, step=None, dtype=None):
|
|
# Allow for improved performance if given literal arguments.
|
|
lit_start = start_value if start_value is not None else start
|
|
lit_stop = stop_value if stop_value is not None else stop
|
|
lit_step = step_value if step_value is not None else step
|
|
|
|
_step = lit_step if lit_step is not None else 1
|
|
if lit_stop is None:
|
|
_start, _stop = 0, lit_start
|
|
else:
|
|
_start, _stop = lit_start, lit_stop
|
|
|
|
if _step == 0:
|
|
raise ValueError("Maximum allowed size exceeded")
|
|
|
|
nitems_c = (_stop - _start) / _step
|
|
nitems_r = int(math.ceil(nitems_c.real))
|
|
|
|
# Binary operator needed for compiler branch pruning.
|
|
if use_complex is True:
|
|
nitems_i = int(math.ceil(nitems_c.imag))
|
|
nitems = max(min(nitems_i, nitems_r), 0)
|
|
else:
|
|
nitems = max(nitems_r, 0)
|
|
arr = np.empty(nitems, true_dtype)
|
|
val = _start
|
|
for i in range(nitems):
|
|
arr[i] = val + (i * _step)
|
|
return arr
|
|
|
|
return impl
|
|
|
|
|
|
@overload(np.linspace)
|
|
def numpy_linspace(start, stop, num=50):
|
|
if not all(isinstance(arg, types.Number) for arg in [start, stop]):
|
|
return
|
|
|
|
if not isinstance(num, (int, types.Integer)):
|
|
msg = 'The argument "num" must be an integer'
|
|
raise errors.TypingError(msg)
|
|
|
|
if any(isinstance(arg, types.Complex) for arg in [start, stop]):
|
|
dtype = types.complex128
|
|
else:
|
|
dtype = types.float64
|
|
|
|
# Implementation based on https://github.com/numpy/numpy/blob/v1.20.0/numpy/core/function_base.py#L24 # noqa: E501
|
|
def linspace(start, stop, num=50):
|
|
arr = np.empty(num, dtype)
|
|
# The multiply by 1.0 mirrors
|
|
# https://github.com/numpy/numpy/blob/v1.20.0/numpy/core/function_base.py#L125-L128 # noqa: E501
|
|
# the side effect of this is important... start and stop become the same
|
|
# type as `dtype` i.e. 64/128 bits wide (float/complex). This is
|
|
# important later when used in the `np.divide`.
|
|
start = start * 1.0
|
|
stop = stop * 1.0
|
|
if num == 0:
|
|
return arr
|
|
div = num - 1
|
|
if div > 0:
|
|
delta = stop - start
|
|
step = np.divide(delta, div)
|
|
for i in range(0, num):
|
|
arr[i] = start + (i * step)
|
|
else:
|
|
arr[0] = start
|
|
if num > 1:
|
|
arr[-1] = stop
|
|
return arr
|
|
return linspace
|
|
|
|
|
|
def _array_copy(context, builder, sig, args):
|
|
"""
|
|
Array copy.
|
|
"""
|
|
arytype = sig.args[0]
|
|
ary = make_array(arytype)(context, builder, value=args[0])
|
|
shapes = cgutils.unpack_tuple(builder, ary.shape)
|
|
|
|
rettype = sig.return_type
|
|
ret = _empty_nd_impl(context, builder, rettype, shapes)
|
|
|
|
src_data = ary.data
|
|
dest_data = ret.data
|
|
|
|
assert rettype.layout in "CF"
|
|
if arytype.layout == rettype.layout:
|
|
# Fast path: memcpy
|
|
cgutils.raw_memcpy(builder, dest_data, src_data, ary.nitems,
|
|
ary.itemsize, align=1)
|
|
|
|
else:
|
|
src_strides = cgutils.unpack_tuple(builder, ary.strides)
|
|
dest_strides = cgutils.unpack_tuple(builder, ret.strides)
|
|
intp_t = context.get_value_type(types.intp)
|
|
|
|
with cgutils.loop_nest(builder, shapes, intp_t) as indices:
|
|
src_ptr = cgutils.get_item_pointer2(context, builder, src_data,
|
|
shapes, src_strides,
|
|
arytype.layout, indices)
|
|
dest_ptr = cgutils.get_item_pointer2(context, builder, dest_data,
|
|
shapes, dest_strides,
|
|
rettype.layout, indices)
|
|
builder.store(builder.load(src_ptr), dest_ptr)
|
|
|
|
return impl_ret_new_ref(context, builder, sig.return_type, ret._getvalue())
|
|
|
|
|
|
@intrinsic
|
|
def _array_copy_intrinsic(typingctx, a):
|
|
assert isinstance(a, types.Array)
|
|
layout = 'F' if a.layout == 'F' else 'C'
|
|
ret = a.copy(layout=layout, readonly=False)
|
|
sig = ret(a)
|
|
return sig, _array_copy
|
|
|
|
|
|
@lower_builtin("array.copy", types.Array)
|
|
def array_copy(context, builder, sig, args):
|
|
return _array_copy(context, builder, sig, args)
|
|
|
|
|
|
@overload(np.copy)
|
|
def impl_numpy_copy(a):
|
|
if isinstance(a, types.Array):
|
|
def numpy_copy(a):
|
|
return _array_copy_intrinsic(a)
|
|
return numpy_copy
|
|
|
|
|
|
def _as_layout_array(context, builder, sig, args, output_layout):
|
|
"""
|
|
Common logic for layout conversion function;
|
|
e.g. ascontiguousarray and asfortranarray
|
|
"""
|
|
retty = sig.return_type
|
|
aryty = sig.args[0]
|
|
assert retty.layout == output_layout, 'return-type has incorrect layout'
|
|
|
|
if aryty.ndim == 0:
|
|
# 0-dim input => asfortranarray() returns a 1-dim array
|
|
assert retty.ndim == 1
|
|
ary = make_array(aryty)(context, builder, value=args[0])
|
|
ret = make_array(retty)(context, builder)
|
|
|
|
shape = context.get_constant_generic(
|
|
builder, types.UniTuple(types.intp, 1), (1,),
|
|
)
|
|
strides = context.make_tuple(builder,
|
|
types.UniTuple(types.intp, 1),
|
|
(ary.itemsize,))
|
|
populate_array(ret, ary.data, shape, strides, ary.itemsize,
|
|
ary.meminfo, ary.parent)
|
|
return impl_ret_borrowed(context, builder, retty, ret._getvalue())
|
|
|
|
elif (retty.layout == aryty.layout
|
|
or (aryty.ndim == 1 and aryty.layout in 'CF')):
|
|
# 1-dim contiguous input => return the same array
|
|
return impl_ret_borrowed(context, builder, retty, args[0])
|
|
|
|
else:
|
|
if aryty.layout == 'A':
|
|
# There's still chance the array is in contiguous layout,
|
|
# just that we don't know at compile time.
|
|
# We can do a runtime check.
|
|
|
|
# Prepare and call is_contiguous or is_fortran
|
|
assert output_layout in 'CF'
|
|
check_func = is_contiguous if output_layout == 'C' else is_fortran
|
|
is_contig = _call_contiguous_check(check_func,
|
|
context,
|
|
builder,
|
|
aryty,
|
|
args[0])
|
|
with builder.if_else(is_contig) as (then, orelse):
|
|
# If the array is already contiguous, just return it
|
|
with then:
|
|
out_then = impl_ret_borrowed(context, builder, retty,
|
|
args[0])
|
|
then_blk = builder.block
|
|
# Otherwise, copy to a new contiguous region
|
|
with orelse:
|
|
out_orelse = _array_copy(context, builder, sig, args)
|
|
orelse_blk = builder.block
|
|
# Phi node for the return value
|
|
ret_phi = builder.phi(out_then.type)
|
|
ret_phi.add_incoming(out_then, then_blk)
|
|
ret_phi.add_incoming(out_orelse, orelse_blk)
|
|
return ret_phi
|
|
|
|
else:
|
|
# Return a copy with the right layout
|
|
return _array_copy(context, builder, sig, args)
|
|
|
|
|
|
@intrinsic
|
|
def _as_layout_array_intrinsic(typingctx, a, output_layout):
|
|
if not isinstance(output_layout, types.StringLiteral):
|
|
raise errors.RequireLiteralValue(output_layout)
|
|
|
|
ret = a.copy(layout=output_layout.literal_value, ndim=max(a.ndim, 1))
|
|
sig = ret(a, output_layout)
|
|
|
|
return sig, lambda c, b, s, a: _as_layout_array(
|
|
c, b, s, a, output_layout=output_layout.literal_value)
|
|
|
|
|
|
@overload(np.ascontiguousarray)
|
|
def array_ascontiguousarray(a):
|
|
if not type_can_asarray(a):
|
|
raise errors.TypingError('The argument "a" must be array-like')
|
|
|
|
if isinstance(a, (types.Number, types.Boolean,)):
|
|
def impl(a):
|
|
return np.ascontiguousarray(np.array(a))
|
|
elif isinstance(a, types.Array):
|
|
def impl(a):
|
|
return _as_layout_array_intrinsic(a, 'C')
|
|
return impl
|
|
|
|
|
|
@overload(np.asfortranarray)
|
|
def array_asfortranarray(a):
|
|
if not type_can_asarray(a):
|
|
raise errors.TypingError('The argument "a" must be array-like')
|
|
|
|
if isinstance(a, (types.Number, types.Boolean,)):
|
|
def impl(a):
|
|
return np.asfortranarray(np.array(a))
|
|
return impl
|
|
elif isinstance(a, types.Array):
|
|
def impl(a):
|
|
return _as_layout_array_intrinsic(a, 'F')
|
|
return impl
|
|
|
|
|
|
@lower_builtin("array.astype", types.Array, types.DTypeSpec)
|
|
@lower_builtin("array.astype", types.Array, types.StringLiteral)
|
|
def array_astype(context, builder, sig, args):
|
|
arytype = sig.args[0]
|
|
ary = make_array(arytype)(context, builder, value=args[0])
|
|
shapes = cgutils.unpack_tuple(builder, ary.shape)
|
|
|
|
rettype = sig.return_type
|
|
ret = _empty_nd_impl(context, builder, rettype, shapes)
|
|
|
|
src_data = ary.data
|
|
dest_data = ret.data
|
|
|
|
src_strides = cgutils.unpack_tuple(builder, ary.strides)
|
|
dest_strides = cgutils.unpack_tuple(builder, ret.strides)
|
|
intp_t = context.get_value_type(types.intp)
|
|
|
|
with cgutils.loop_nest(builder, shapes, intp_t) as indices:
|
|
src_ptr = cgutils.get_item_pointer2(context, builder, src_data,
|
|
shapes, src_strides,
|
|
arytype.layout, indices)
|
|
dest_ptr = cgutils.get_item_pointer2(context, builder, dest_data,
|
|
shapes, dest_strides,
|
|
rettype.layout, indices)
|
|
item = load_item(context, builder, arytype, src_ptr)
|
|
item = context.cast(builder, item, arytype.dtype, rettype.dtype)
|
|
store_item(context, builder, rettype, item, dest_ptr)
|
|
|
|
return impl_ret_new_ref(context, builder, sig.return_type, ret._getvalue())
|
|
|
|
|
|
@intrinsic
|
|
def np_frombuffer(typingctx, buffer, dtype, retty):
|
|
ty = retty.instance_type
|
|
sig = ty(buffer, dtype, retty)
|
|
|
|
def codegen(context, builder, sig, args):
|
|
bufty = sig.args[0]
|
|
aryty = sig.return_type
|
|
|
|
buf = make_array(bufty)(context, builder, value=args[0])
|
|
out_ary_ty = make_array(aryty)
|
|
out_ary = out_ary_ty(context, builder)
|
|
out_datamodel = out_ary._datamodel
|
|
|
|
itemsize = get_itemsize(context, aryty)
|
|
ll_itemsize = Constant(buf.itemsize.type, itemsize)
|
|
nbytes = builder.mul(buf.nitems, buf.itemsize)
|
|
|
|
# Check that the buffer size is compatible
|
|
rem = builder.srem(nbytes, ll_itemsize)
|
|
is_incompatible = cgutils.is_not_null(builder, rem)
|
|
with builder.if_then(is_incompatible, likely=False):
|
|
msg = "buffer size must be a multiple of element size"
|
|
context.call_conv.return_user_exc(builder, ValueError, (msg,))
|
|
|
|
shape = cgutils.pack_array(builder, [builder.sdiv(nbytes, ll_itemsize)])
|
|
strides = cgutils.pack_array(builder, [ll_itemsize])
|
|
data = builder.bitcast(
|
|
buf.data, context.get_value_type(out_datamodel.get_type('data'))
|
|
)
|
|
|
|
populate_array(out_ary,
|
|
data=data,
|
|
shape=shape,
|
|
strides=strides,
|
|
itemsize=ll_itemsize,
|
|
meminfo=buf.meminfo,
|
|
parent=buf.parent,)
|
|
|
|
res = out_ary._getvalue()
|
|
return impl_ret_borrowed(context, builder, sig.return_type, res)
|
|
return sig, codegen
|
|
|
|
|
|
@overload(np.frombuffer)
|
|
def impl_np_frombuffer(buffer, dtype=float):
|
|
_check_const_str_dtype("frombuffer", dtype)
|
|
|
|
if not isinstance(buffer, types.Buffer) or buffer.layout != 'C':
|
|
msg = f'Argument "buffer" must be buffer-like. Got {buffer}'
|
|
raise errors.TypingError(msg)
|
|
|
|
if (dtype is float or
|
|
(isinstance(dtype, types.Function) and dtype.typing_key is float) or
|
|
is_nonelike(dtype)): #default
|
|
nb_dtype = types.double
|
|
else:
|
|
nb_dtype = ty_parse_dtype(dtype)
|
|
|
|
if nb_dtype is not None:
|
|
retty = types.Array(dtype=nb_dtype, ndim=1, layout='C',
|
|
readonly=not buffer.mutable)
|
|
else:
|
|
msg = ("Cannot parse input types to function "
|
|
f"np.frombuffer({buffer}, {dtype})")
|
|
raise errors.TypingError(msg)
|
|
|
|
def impl(buffer, dtype=float):
|
|
return np_frombuffer(buffer, dtype, retty)
|
|
return impl
|
|
|
|
|
|
@overload(carray)
|
|
def impl_carray(ptr, shape, dtype=None):
|
|
if is_nonelike(dtype):
|
|
intrinsic_cfarray = get_cfarray_intrinsic('C', None)
|
|
|
|
def impl(ptr, shape, dtype=None):
|
|
return intrinsic_cfarray(ptr, shape)
|
|
return impl
|
|
elif isinstance(dtype, types.DTypeSpec):
|
|
intrinsic_cfarray = get_cfarray_intrinsic('C', dtype)
|
|
|
|
def impl(ptr, shape, dtype=None):
|
|
return intrinsic_cfarray(ptr, shape)
|
|
return impl
|
|
|
|
|
|
@overload(farray)
|
|
def impl_farray(ptr, shape, dtype=None):
|
|
if is_nonelike(dtype):
|
|
intrinsic_cfarray = get_cfarray_intrinsic('F', None)
|
|
|
|
def impl(ptr, shape, dtype=None):
|
|
return intrinsic_cfarray(ptr, shape)
|
|
return impl
|
|
elif isinstance(dtype, types.DTypeSpec):
|
|
intrinsic_cfarray = get_cfarray_intrinsic('F', dtype)
|
|
|
|
def impl(ptr, shape, dtype=None):
|
|
return intrinsic_cfarray(ptr, shape)
|
|
return impl
|
|
|
|
|
|
def get_cfarray_intrinsic(layout, dtype_):
|
|
@intrinsic
|
|
def intrinsic_cfarray(typingctx, ptr, shape):
|
|
if ptr is types.voidptr:
|
|
ptr_dtype = None
|
|
elif isinstance(ptr, types.CPointer):
|
|
ptr_dtype = ptr.dtype
|
|
else:
|
|
msg = f"pointer argument expected, got '{ptr}'"
|
|
raise errors.NumbaTypeError(msg)
|
|
|
|
if dtype_ is None:
|
|
if ptr_dtype is None:
|
|
msg = "explicit dtype required for void* argument"
|
|
raise errors.NumbaTypeError(msg)
|
|
dtype = ptr_dtype
|
|
elif isinstance(dtype_, types.DTypeSpec):
|
|
dtype = dtype_.dtype
|
|
if ptr_dtype is not None and dtype != ptr_dtype:
|
|
msg = f"mismatching dtype '{dtype}' for pointer type '{ptr}'"
|
|
raise errors.NumbaTypeError(msg)
|
|
else:
|
|
msg = f"invalid dtype spec '{dtype_}'"
|
|
raise errors.NumbaTypeError(msg)
|
|
|
|
ndim = ty_parse_shape(shape)
|
|
if ndim is None:
|
|
msg = f"invalid shape '{shape}'"
|
|
raise errors.NumbaTypeError(msg)
|
|
|
|
retty = types.Array(dtype, ndim, layout)
|
|
sig = signature(retty, ptr, shape)
|
|
return sig, np_cfarray
|
|
return intrinsic_cfarray
|
|
|
|
|
|
def np_cfarray(context, builder, sig, args):
|
|
"""
|
|
numba.numpy_support.carray(...) and
|
|
numba.numpy_support.farray(...).
|
|
"""
|
|
ptrty, shapety = sig.args[:2]
|
|
ptr, shape = args[:2]
|
|
|
|
aryty = sig.return_type
|
|
assert aryty.layout in 'CF'
|
|
|
|
out_ary = make_array(aryty)(context, builder)
|
|
|
|
itemsize = get_itemsize(context, aryty)
|
|
ll_itemsize = cgutils.intp_t(itemsize)
|
|
|
|
if isinstance(shapety, types.BaseTuple):
|
|
shapes = cgutils.unpack_tuple(builder, shape)
|
|
else:
|
|
shapety = (shapety,)
|
|
shapes = (shape,)
|
|
shapes = [context.cast(builder, value, fromty, types.intp)
|
|
for fromty, value in zip(shapety, shapes)]
|
|
|
|
off = ll_itemsize
|
|
strides = []
|
|
if aryty.layout == 'F':
|
|
for s in shapes:
|
|
strides.append(off)
|
|
off = builder.mul(off, s)
|
|
else:
|
|
for s in reversed(shapes):
|
|
strides.append(off)
|
|
off = builder.mul(off, s)
|
|
strides.reverse()
|
|
|
|
data = builder.bitcast(ptr,
|
|
context.get_data_type(aryty.dtype).as_pointer())
|
|
|
|
populate_array(out_ary,
|
|
data=data,
|
|
shape=shapes,
|
|
strides=strides,
|
|
itemsize=ll_itemsize,
|
|
# Array is not memory-managed
|
|
meminfo=None,
|
|
)
|
|
|
|
res = out_ary._getvalue()
|
|
return impl_ret_new_ref(context, builder, sig.return_type, res)
|
|
|
|
|
|
def _get_seq_size(context, builder, seqty, seq):
|
|
if isinstance(seqty, types.BaseTuple):
|
|
return context.get_constant(types.intp, len(seqty))
|
|
elif isinstance(seqty, types.Sequence):
|
|
len_impl = context.get_function(len, signature(types.intp, seqty,))
|
|
return len_impl(builder, (seq,))
|
|
else:
|
|
assert 0
|
|
|
|
|
|
def _get_borrowing_getitem(context, seqty):
|
|
"""
|
|
Return a getitem() implementation that doesn't incref its result.
|
|
"""
|
|
retty = seqty.dtype
|
|
getitem_impl = context.get_function(operator.getitem,
|
|
signature(retty, seqty, types.intp))
|
|
|
|
def wrap(builder, args):
|
|
ret = getitem_impl(builder, args)
|
|
if context.enable_nrt:
|
|
context.nrt.decref(builder, retty, ret)
|
|
return ret
|
|
|
|
return wrap
|
|
|
|
|
|
def compute_sequence_shape(context, builder, ndim, seqty, seq):
|
|
"""
|
|
Compute the likely shape of a nested sequence (possibly 0d).
|
|
"""
|
|
intp_t = context.get_value_type(types.intp)
|
|
zero = Constant(intp_t, 0)
|
|
|
|
def get_first_item(seqty, seq):
|
|
if isinstance(seqty, types.BaseTuple):
|
|
if len(seqty) == 0:
|
|
return None, None
|
|
else:
|
|
return seqty[0], builder.extract_value(seq, 0)
|
|
else:
|
|
getitem_impl = _get_borrowing_getitem(context, seqty)
|
|
return seqty.dtype, getitem_impl(builder, (seq, zero))
|
|
|
|
# Compute shape by traversing the first element of each nested
|
|
# sequence
|
|
shapes = []
|
|
innerty, inner = seqty, seq
|
|
|
|
for i in range(ndim):
|
|
if i > 0:
|
|
innerty, inner = get_first_item(innerty, inner)
|
|
shapes.append(_get_seq_size(context, builder, innerty, inner))
|
|
|
|
return tuple(shapes)
|
|
|
|
|
|
def check_sequence_shape(context, builder, seqty, seq, shapes):
|
|
"""
|
|
Check the nested sequence matches the given *shapes*.
|
|
"""
|
|
|
|
def _fail():
|
|
context.call_conv.return_user_exc(builder, ValueError,
|
|
("incompatible sequence shape",))
|
|
|
|
def check_seq_size(seqty, seq, shapes):
|
|
if len(shapes) == 0:
|
|
return
|
|
|
|
size = _get_seq_size(context, builder, seqty, seq)
|
|
expected = shapes[0]
|
|
mismatch = builder.icmp_signed('!=', size, expected)
|
|
with builder.if_then(mismatch, likely=False):
|
|
_fail()
|
|
|
|
if len(shapes) == 1:
|
|
return
|
|
|
|
if isinstance(seqty, types.Sequence):
|
|
getitem_impl = _get_borrowing_getitem(context, seqty)
|
|
with cgutils.for_range(builder, size) as loop:
|
|
innerty = seqty.dtype
|
|
inner = getitem_impl(builder, (seq, loop.index))
|
|
check_seq_size(innerty, inner, shapes[1:])
|
|
|
|
elif isinstance(seqty, types.BaseTuple):
|
|
for i in range(len(seqty)):
|
|
innerty = seqty[i]
|
|
inner = builder.extract_value(seq, i)
|
|
check_seq_size(innerty, inner, shapes[1:])
|
|
|
|
else:
|
|
assert 0, seqty
|
|
|
|
check_seq_size(seqty, seq, shapes)
|
|
|
|
|
|
def assign_sequence_to_array(context, builder, data, shapes, strides,
|
|
arrty, seqty, seq):
|
|
"""
|
|
Assign a nested sequence contents to an array. The shape must match
|
|
the sequence's structure.
|
|
"""
|
|
|
|
def assign_item(indices, valty, val):
|
|
ptr = cgutils.get_item_pointer2(context, builder, data, shapes, strides,
|
|
arrty.layout, indices, wraparound=False)
|
|
val = context.cast(builder, val, valty, arrty.dtype)
|
|
store_item(context, builder, arrty, val, ptr)
|
|
|
|
def assign(seqty, seq, shapes, indices):
|
|
if len(shapes) == 0:
|
|
assert not isinstance(seqty, (types.Sequence, types.BaseTuple))
|
|
assign_item(indices, seqty, seq)
|
|
return
|
|
|
|
size = shapes[0]
|
|
|
|
if isinstance(seqty, types.Sequence):
|
|
getitem_impl = _get_borrowing_getitem(context, seqty)
|
|
with cgutils.for_range(builder, size) as loop:
|
|
innerty = seqty.dtype
|
|
inner = getitem_impl(builder, (seq, loop.index))
|
|
assign(innerty, inner, shapes[1:], indices + (loop.index,))
|
|
|
|
elif isinstance(seqty, types.BaseTuple):
|
|
for i in range(len(seqty)):
|
|
innerty = seqty[i]
|
|
inner = builder.extract_value(seq, i)
|
|
index = context.get_constant(types.intp, i)
|
|
assign(innerty, inner, shapes[1:], indices + (index,))
|
|
|
|
else:
|
|
assert 0, seqty
|
|
|
|
assign(seqty, seq, shapes, ())
|
|
|
|
|
|
def np_array_typer(typingctx, object, dtype):
|
|
ndim, seq_dtype = _parse_nested_sequence(typingctx, object)
|
|
if is_nonelike(dtype):
|
|
dtype = seq_dtype
|
|
else:
|
|
dtype = ty_parse_dtype(dtype)
|
|
if dtype is None:
|
|
return
|
|
return types.Array(dtype, ndim, 'C')
|
|
|
|
|
|
@intrinsic
|
|
def np_array(typingctx, obj, dtype):
|
|
_check_const_str_dtype("array", dtype)
|
|
ret = np_array_typer(typingctx, obj, dtype)
|
|
sig = ret(obj, dtype)
|
|
|
|
def codegen(context, builder, sig, args):
|
|
arrty = sig.return_type
|
|
ndim = arrty.ndim
|
|
seqty = sig.args[0]
|
|
seq = args[0]
|
|
|
|
shapes = compute_sequence_shape(context, builder, ndim, seqty, seq)
|
|
assert len(shapes) == ndim
|
|
|
|
check_sequence_shape(context, builder, seqty, seq, shapes)
|
|
arr = _empty_nd_impl(context, builder, arrty, shapes)
|
|
assign_sequence_to_array(context, builder, arr.data, shapes,
|
|
arr.strides, arrty, seqty, seq)
|
|
|
|
return impl_ret_new_ref(context, builder, sig.return_type,
|
|
arr._getvalue())
|
|
|
|
return sig, codegen
|
|
|
|
|
|
@overload(np.array)
|
|
def impl_np_array(object, dtype=None):
|
|
_check_const_str_dtype("array", dtype)
|
|
if not type_can_asarray(object):
|
|
raise errors.TypingError('The argument "object" must '
|
|
'be array-like')
|
|
if not is_nonelike(dtype) and ty_parse_dtype(dtype) is None:
|
|
msg = 'The argument "dtype" must be a data-type if it is provided'
|
|
raise errors.TypingError(msg)
|
|
|
|
def impl(object, dtype=None):
|
|
return np_array(object, dtype)
|
|
return impl
|
|
|
|
|
|
def _normalize_axis(context, builder, func_name, ndim, axis):
|
|
zero = axis.type(0)
|
|
ll_ndim = axis.type(ndim)
|
|
|
|
# Normalize negative axis
|
|
is_neg_axis = builder.icmp_signed('<', axis, zero)
|
|
axis = builder.select(is_neg_axis, builder.add(axis, ll_ndim), axis)
|
|
|
|
# Check axis for bounds
|
|
axis_out_of_bounds = builder.or_(
|
|
builder.icmp_signed('<', axis, zero),
|
|
builder.icmp_signed('>=', axis, ll_ndim))
|
|
with builder.if_then(axis_out_of_bounds, likely=False):
|
|
msg = "%s(): axis out of bounds" % func_name
|
|
context.call_conv.return_user_exc(builder, IndexError, (msg,))
|
|
|
|
return axis
|
|
|
|
|
|
def _insert_axis_in_shape(context, builder, orig_shape, ndim, axis):
|
|
"""
|
|
Compute shape with the new axis inserted
|
|
e.g. given original shape (2, 3, 4) and axis=2,
|
|
the returned new shape is (2, 3, 1, 4).
|
|
"""
|
|
assert len(orig_shape) == ndim - 1
|
|
|
|
ll_shty = ir.ArrayType(cgutils.intp_t, ndim)
|
|
shapes = cgutils.alloca_once(builder, ll_shty)
|
|
|
|
one = cgutils.intp_t(1)
|
|
|
|
# 1. copy original sizes at appropriate places
|
|
for dim in range(ndim - 1):
|
|
ll_dim = cgutils.intp_t(dim)
|
|
after_axis = builder.icmp_signed('>=', ll_dim, axis)
|
|
sh = orig_shape[dim]
|
|
idx = builder.select(after_axis,
|
|
builder.add(ll_dim, one),
|
|
ll_dim)
|
|
builder.store(sh, cgutils.gep_inbounds(builder, shapes, 0, idx))
|
|
|
|
# 2. insert new size (1) at axis dimension
|
|
builder.store(one, cgutils.gep_inbounds(builder, shapes, 0, axis))
|
|
|
|
return cgutils.unpack_tuple(builder, builder.load(shapes))
|
|
|
|
|
|
def _insert_axis_in_strides(context, builder, orig_strides, ndim, axis):
|
|
"""
|
|
Same as _insert_axis_in_shape(), but with a strides array.
|
|
"""
|
|
assert len(orig_strides) == ndim - 1
|
|
|
|
ll_shty = ir.ArrayType(cgutils.intp_t, ndim)
|
|
strides = cgutils.alloca_once(builder, ll_shty)
|
|
|
|
one = cgutils.intp_t(1)
|
|
zero = cgutils.intp_t(0)
|
|
|
|
# 1. copy original strides at appropriate places
|
|
for dim in range(ndim - 1):
|
|
ll_dim = cgutils.intp_t(dim)
|
|
after_axis = builder.icmp_signed('>=', ll_dim, axis)
|
|
idx = builder.select(after_axis,
|
|
builder.add(ll_dim, one),
|
|
ll_dim)
|
|
builder.store(orig_strides[dim],
|
|
cgutils.gep_inbounds(builder, strides, 0, idx))
|
|
|
|
# 2. insert new stride at axis dimension
|
|
# (the value is indifferent for a 1-sized dimension, we use 0)
|
|
builder.store(zero, cgutils.gep_inbounds(builder, strides, 0, axis))
|
|
|
|
return cgutils.unpack_tuple(builder, builder.load(strides))
|
|
|
|
|
|
def expand_dims(context, builder, sig, args, axis):
|
|
"""
|
|
np.expand_dims() with the given axis.
|
|
"""
|
|
retty = sig.return_type
|
|
ndim = retty.ndim
|
|
arrty = sig.args[0]
|
|
|
|
arr = make_array(arrty)(context, builder, value=args[0])
|
|
ret = make_array(retty)(context, builder)
|
|
|
|
shapes = cgutils.unpack_tuple(builder, arr.shape)
|
|
strides = cgutils.unpack_tuple(builder, arr.strides)
|
|
|
|
new_shapes = _insert_axis_in_shape(context, builder, shapes, ndim, axis)
|
|
new_strides = _insert_axis_in_strides(context, builder, strides, ndim, axis)
|
|
|
|
populate_array(ret,
|
|
data=arr.data,
|
|
shape=new_shapes,
|
|
strides=new_strides,
|
|
itemsize=arr.itemsize,
|
|
meminfo=arr.meminfo,
|
|
parent=arr.parent)
|
|
|
|
return ret._getvalue()
|
|
|
|
|
|
@intrinsic
|
|
def np_expand_dims(typingctx, a, axis):
|
|
layout = a.layout if a.ndim <= 1 else 'A'
|
|
ret = a.copy(ndim=a.ndim + 1, layout=layout)
|
|
sig = ret(a, axis)
|
|
|
|
def codegen(context, builder, sig, args):
|
|
axis = context.cast(builder, args[1], sig.args[1], types.intp)
|
|
axis = _normalize_axis(context, builder, "np.expand_dims",
|
|
sig.return_type.ndim, axis)
|
|
|
|
ret = expand_dims(context, builder, sig, args, axis)
|
|
return impl_ret_borrowed(context, builder, sig.return_type, ret)
|
|
|
|
return sig, codegen
|
|
|
|
|
|
@overload(np.expand_dims)
|
|
def impl_np_expand_dims(a, axis):
|
|
if not isinstance(a, types.Array):
|
|
msg = f'First argument "a" must be an array. Got {a}'
|
|
raise errors.TypingError(msg)
|
|
|
|
if not isinstance(axis, types.Integer):
|
|
msg = f'Argument "axis" must be an integer. Got {axis}'
|
|
raise errors.TypingError(msg)
|
|
|
|
def impl(a, axis):
|
|
return np_expand_dims(a, axis)
|
|
return impl
|
|
|
|
|
|
def _atleast_nd(minimum, axes):
|
|
@intrinsic
|
|
def impl(typingcontext, *args):
|
|
arrtys = args
|
|
rettys = [arg.copy(ndim=max(arg.ndim, minimum)) for arg in args]
|
|
|
|
def codegen(context, builder, sig, args):
|
|
transform = _atleast_nd_transform(minimum, axes)
|
|
arrs = cgutils.unpack_tuple(builder, args[0])
|
|
|
|
rets = [transform(context, builder, arr, arrty, retty)
|
|
for arr, arrty, retty in zip(arrs, arrtys, rettys)]
|
|
|
|
if len(rets) > 1:
|
|
ret = context.make_tuple(builder, sig.return_type, rets)
|
|
else:
|
|
ret = rets[0]
|
|
return impl_ret_borrowed(context, builder, sig.return_type, ret)
|
|
|
|
return signature(types.Tuple(rettys) if len(rettys) > 1 else rettys[0],
|
|
types.StarArgTuple.from_types(args)), codegen
|
|
|
|
return lambda *args: impl(*args)
|
|
|
|
|
|
def _atleast_nd_transform(min_ndim, axes):
|
|
"""
|
|
Return a callback successively inserting 1-sized dimensions at the
|
|
following axes.
|
|
"""
|
|
assert min_ndim == len(axes)
|
|
|
|
def transform(context, builder, arr, arrty, retty):
|
|
for i in range(min_ndim):
|
|
ndim = i + 1
|
|
if arrty.ndim < ndim:
|
|
axis = cgutils.intp_t(axes[i])
|
|
newarrty = arrty.copy(ndim=arrty.ndim + 1)
|
|
arr = expand_dims(context, builder,
|
|
typing.signature(newarrty, arrty), (arr,),
|
|
axis)
|
|
arrty = newarrty
|
|
|
|
return arr
|
|
|
|
return transform
|
|
|
|
|
|
@overload(np.atleast_1d)
|
|
def np_atleast_1d(*args):
|
|
if all(isinstance(arg, types.Array) for arg in args):
|
|
return _atleast_nd(1, [0])
|
|
|
|
|
|
@overload(np.atleast_2d)
|
|
def np_atleast_2d(*args):
|
|
if all(isinstance(arg, types.Array) for arg in args):
|
|
return _atleast_nd(2, [0, 0])
|
|
|
|
|
|
@overload(np.atleast_3d)
|
|
def np_atleast_3d(*args):
|
|
if all(isinstance(arg, types.Array) for arg in args):
|
|
return _atleast_nd(3, [0, 0, 2])
|
|
|
|
|
|
def _do_concatenate(context, builder, axis,
|
|
arrtys, arrs, arr_shapes, arr_strides,
|
|
retty, ret_shapes):
|
|
"""
|
|
Concatenate arrays along the given axis.
|
|
"""
|
|
assert len(arrtys) == len(arrs) == len(arr_shapes) == len(arr_strides)
|
|
|
|
zero = cgutils.intp_t(0)
|
|
|
|
# Allocate return array
|
|
ret = _empty_nd_impl(context, builder, retty, ret_shapes)
|
|
ret_strides = cgutils.unpack_tuple(builder, ret.strides)
|
|
|
|
# Compute the offset by which to bump the destination pointer
|
|
# after copying each input array.
|
|
# Morally, we need to copy each input array at different start indices
|
|
# into the destination array; bumping the destination pointer
|
|
# is simply easier than offsetting all destination indices.
|
|
copy_offsets = []
|
|
|
|
for arr_sh in arr_shapes:
|
|
# offset = ret_strides[axis] * input_shape[axis]
|
|
offset = zero
|
|
for dim, (size, stride) in enumerate(zip(arr_sh, ret_strides)):
|
|
is_axis = builder.icmp_signed('==', axis.type(dim), axis)
|
|
addend = builder.mul(size, stride)
|
|
offset = builder.select(is_axis,
|
|
builder.add(offset, addend),
|
|
offset)
|
|
copy_offsets.append(offset)
|
|
|
|
# Copy input arrays into the return array
|
|
ret_data = ret.data
|
|
|
|
for arrty, arr, arr_sh, arr_st, offset in zip(arrtys, arrs, arr_shapes,
|
|
arr_strides, copy_offsets):
|
|
arr_data = arr.data
|
|
|
|
# Do the copy loop
|
|
# Note the loop nesting is optimized for the destination layout
|
|
loop_nest = cgutils.loop_nest(builder, arr_sh, cgutils.intp_t,
|
|
order=retty.layout)
|
|
|
|
with loop_nest as indices:
|
|
src_ptr = cgutils.get_item_pointer2(context, builder, arr_data,
|
|
arr_sh, arr_st,
|
|
arrty.layout, indices)
|
|
val = load_item(context, builder, arrty, src_ptr)
|
|
val = context.cast(builder, val, arrty.dtype, retty.dtype)
|
|
dest_ptr = cgutils.get_item_pointer2(context, builder, ret_data,
|
|
ret_shapes, ret_strides,
|
|
retty.layout, indices)
|
|
store_item(context, builder, retty, val, dest_ptr)
|
|
|
|
# Bump destination pointer
|
|
ret_data = cgutils.pointer_add(builder, ret_data, offset)
|
|
|
|
return ret
|
|
|
|
|
|
def _np_concatenate(context, builder, arrtys, arrs, retty, axis):
|
|
ndim = retty.ndim
|
|
|
|
arrs = [make_array(aty)(context, builder, value=a)
|
|
for aty, a in zip(arrtys, arrs)]
|
|
|
|
axis = _normalize_axis(context, builder, "np.concatenate", ndim, axis)
|
|
|
|
# Get input shapes
|
|
arr_shapes = [cgutils.unpack_tuple(builder, arr.shape) for arr in arrs]
|
|
arr_strides = [cgutils.unpack_tuple(builder, arr.strides) for arr in arrs]
|
|
|
|
# Compute return shape:
|
|
# - the dimension for the concatenation axis is summed over all inputs
|
|
# - other dimensions must match exactly for each input
|
|
ret_shapes = [cgutils.alloca_once_value(builder, sh)
|
|
for sh in arr_shapes[0]]
|
|
|
|
for dim in range(ndim):
|
|
is_axis = builder.icmp_signed('==', axis.type(dim), axis)
|
|
ret_shape_ptr = ret_shapes[dim]
|
|
ret_sh = builder.load(ret_shape_ptr)
|
|
other_shapes = [sh[dim] for sh in arr_shapes[1:]]
|
|
|
|
with builder.if_else(is_axis) as (on_axis, on_other_dim):
|
|
with on_axis:
|
|
sh = functools.reduce(
|
|
builder.add,
|
|
other_shapes + [ret_sh])
|
|
builder.store(sh, ret_shape_ptr)
|
|
|
|
with on_other_dim:
|
|
is_ok = cgutils.true_bit
|
|
for sh in other_shapes:
|
|
is_ok = builder.and_(is_ok,
|
|
builder.icmp_signed('==', sh, ret_sh))
|
|
with builder.if_then(builder.not_(is_ok), likely=False):
|
|
context.call_conv.return_user_exc(
|
|
builder, ValueError,
|
|
("np.concatenate(): input sizes over "
|
|
"dimension %d do not match" % dim,))
|
|
|
|
ret_shapes = [builder.load(sh) for sh in ret_shapes]
|
|
|
|
ret = _do_concatenate(context, builder, axis,
|
|
arrtys, arrs, arr_shapes, arr_strides,
|
|
retty, ret_shapes)
|
|
return impl_ret_new_ref(context, builder, retty, ret._getvalue())
|
|
|
|
|
|
def _np_stack(context, builder, arrtys, arrs, retty, axis):
|
|
ndim = retty.ndim
|
|
|
|
zero = cgutils.intp_t(0)
|
|
one = cgutils.intp_t(1)
|
|
ll_narrays = cgutils.intp_t(len(arrs))
|
|
|
|
arrs = [make_array(aty)(context, builder, value=a)
|
|
for aty, a in zip(arrtys, arrs)]
|
|
|
|
axis = _normalize_axis(context, builder, "np.stack", ndim, axis)
|
|
|
|
# Check input arrays have the same shape
|
|
orig_shape = cgutils.unpack_tuple(builder, arrs[0].shape)
|
|
|
|
for arr in arrs[1:]:
|
|
is_ok = cgutils.true_bit
|
|
for sh, orig_sh in zip(cgutils.unpack_tuple(builder, arr.shape),
|
|
orig_shape):
|
|
is_ok = builder.and_(is_ok, builder.icmp_signed('==', sh, orig_sh))
|
|
with builder.if_then(builder.not_(is_ok), likely=False):
|
|
context.call_conv.return_user_exc(
|
|
builder, ValueError,
|
|
("np.stack(): all input arrays must have the same shape",))
|
|
|
|
orig_strides = [cgutils.unpack_tuple(builder, arr.strides) for arr in arrs]
|
|
|
|
# Compute input shapes and return shape with the new axis inserted
|
|
# e.g. given 5 input arrays of shape (2, 3, 4) and axis=1,
|
|
# corrected input shape is (2, 1, 3, 4) and return shape is (2, 5, 3, 4).
|
|
ll_shty = ir.ArrayType(cgutils.intp_t, ndim)
|
|
|
|
input_shapes = cgutils.alloca_once(builder, ll_shty)
|
|
ret_shapes = cgutils.alloca_once(builder, ll_shty)
|
|
|
|
# 1. copy original sizes at appropriate places
|
|
for dim in range(ndim - 1):
|
|
ll_dim = cgutils.intp_t(dim)
|
|
after_axis = builder.icmp_signed('>=', ll_dim, axis)
|
|
sh = orig_shape[dim]
|
|
idx = builder.select(after_axis,
|
|
builder.add(ll_dim, one),
|
|
ll_dim)
|
|
builder.store(sh, cgutils.gep_inbounds(builder, input_shapes, 0, idx))
|
|
builder.store(sh, cgutils.gep_inbounds(builder, ret_shapes, 0, idx))
|
|
|
|
# 2. insert new size at axis dimension
|
|
builder.store(one, cgutils.gep_inbounds(builder, input_shapes, 0, axis))
|
|
builder.store(ll_narrays, cgutils.gep_inbounds(builder,
|
|
ret_shapes,
|
|
0,
|
|
axis))
|
|
|
|
input_shapes = cgutils.unpack_tuple(builder, builder.load(input_shapes))
|
|
input_shapes = [input_shapes] * len(arrs)
|
|
ret_shapes = cgutils.unpack_tuple(builder, builder.load(ret_shapes))
|
|
|
|
# Compute input strides for each array with the new axis inserted
|
|
input_strides = [cgutils.alloca_once(builder, ll_shty)
|
|
for i in range(len(arrs))]
|
|
|
|
# 1. copy original strides at appropriate places
|
|
for dim in range(ndim - 1):
|
|
ll_dim = cgutils.intp_t(dim)
|
|
after_axis = builder.icmp_signed('>=', ll_dim, axis)
|
|
idx = builder.select(after_axis,
|
|
builder.add(ll_dim, one),
|
|
ll_dim)
|
|
for i in range(len(arrs)):
|
|
builder.store(orig_strides[i][dim],
|
|
cgutils.gep_inbounds(builder, input_strides[i], 0,
|
|
idx))
|
|
|
|
# 2. insert new stride at axis dimension
|
|
# (the value is indifferent for a 1-sized dimension, we put 0)
|
|
for i in range(len(arrs)):
|
|
builder.store(zero, cgutils.gep_inbounds(builder, input_strides[i], 0,
|
|
axis))
|
|
|
|
input_strides = [cgutils.unpack_tuple(builder, builder.load(st))
|
|
for st in input_strides]
|
|
|
|
# Create concatenated array
|
|
ret = _do_concatenate(context, builder, axis,
|
|
arrtys, arrs, input_shapes, input_strides,
|
|
retty, ret_shapes)
|
|
return impl_ret_new_ref(context, builder, retty, ret._getvalue())
|
|
|
|
|
|
def np_concatenate_typer(typingctx, arrays, axis):
|
|
if axis is not None and not isinstance(axis, types.Integer):
|
|
# Note Numpy allows axis=None, but it isn't documented:
|
|
# https://github.com/numpy/numpy/issues/7968
|
|
return
|
|
|
|
# does type checking
|
|
dtype, ndim = _sequence_of_arrays(typingctx,
|
|
"np.concatenate", arrays)
|
|
if ndim == 0:
|
|
raise TypeError("zero-dimensional arrays cannot be concatenated")
|
|
|
|
layout = _choose_concatenation_layout(arrays)
|
|
|
|
return types.Array(dtype, ndim, layout)
|
|
|
|
|
|
@intrinsic
|
|
def np_concatenate(typingctx, arrays, axis):
|
|
ret = np_concatenate_typer(typingctx, arrays, axis)
|
|
assert isinstance(ret, types.Array)
|
|
sig = ret(arrays, axis)
|
|
|
|
def codegen(context, builder, sig, args):
|
|
axis = context.cast(builder, args[1], sig.args[1], types.intp)
|
|
return _np_concatenate(context, builder,
|
|
list(sig.args[0]),
|
|
cgutils.unpack_tuple(builder, args[0]),
|
|
sig.return_type,
|
|
axis)
|
|
|
|
return sig, codegen
|
|
|
|
|
|
@overload(np.concatenate)
|
|
def impl_np_concatenate(arrays, axis=0):
|
|
if isinstance(arrays, types.BaseTuple):
|
|
def impl(arrays, axis=0):
|
|
return np_concatenate(arrays, axis)
|
|
return impl
|
|
|
|
|
|
def _column_stack_dims(context, func_name, arrays):
|
|
# column_stack() allows stacking 1-d and 2-d arrays together
|
|
for a in arrays:
|
|
if a.ndim < 1 or a.ndim > 2:
|
|
raise TypeError("np.column_stack() is only defined on "
|
|
"1-d and 2-d arrays")
|
|
return 2
|
|
|
|
|
|
@intrinsic
|
|
def np_column_stack(typingctx, tup):
|
|
dtype, ndim = _sequence_of_arrays(typingctx,
|
|
"np.column_stack", tup,
|
|
dim_chooser=_column_stack_dims)
|
|
layout = _choose_concatenation_layout(tup)
|
|
ret = types.Array(dtype, ndim, layout)
|
|
sig = ret(tup)
|
|
|
|
def codegen(context, builder, sig, args):
|
|
orig_arrtys = list(sig.args[0])
|
|
orig_arrs = cgutils.unpack_tuple(builder, args[0])
|
|
|
|
arrtys = []
|
|
arrs = []
|
|
|
|
axis = context.get_constant(types.intp, 1)
|
|
|
|
for arrty, arr in zip(orig_arrtys, orig_arrs):
|
|
if arrty.ndim == 2:
|
|
arrtys.append(arrty)
|
|
arrs.append(arr)
|
|
else:
|
|
# Convert 1d array to 2d column array: np.expand_dims(a, 1)
|
|
assert arrty.ndim == 1
|
|
newty = arrty.copy(ndim=2)
|
|
expand_sig = typing.signature(newty, arrty)
|
|
newarr = expand_dims(context, builder, expand_sig, (arr,), axis)
|
|
|
|
arrtys.append(newty)
|
|
arrs.append(newarr)
|
|
|
|
return _np_concatenate(context, builder, arrtys, arrs,
|
|
sig.return_type, axis)
|
|
|
|
return sig, codegen
|
|
|
|
|
|
@overload(np.column_stack)
|
|
def impl_column_stack(tup):
|
|
if isinstance(tup, types.BaseTuple):
|
|
def impl(tup):
|
|
return np_column_stack(tup)
|
|
return impl
|
|
|
|
|
|
def _np_stack_common(context, builder, sig, args, axis):
|
|
"""
|
|
np.stack() with the given axis value.
|
|
"""
|
|
return _np_stack(context, builder,
|
|
list(sig.args[0]),
|
|
cgutils.unpack_tuple(builder, args[0]),
|
|
sig.return_type,
|
|
axis)
|
|
|
|
|
|
@intrinsic
|
|
def np_stack_common(typingctx, arrays, axis):
|
|
# does type checking
|
|
dtype, ndim = _sequence_of_arrays(typingctx,
|
|
"np.stack", arrays)
|
|
layout = 'F' if all(a.layout == 'F' for a in arrays) else 'C'
|
|
ret = types.Array(dtype, ndim + 1, layout)
|
|
sig = ret(arrays, axis)
|
|
|
|
def codegen(context, builder, sig, args):
|
|
axis = context.cast(builder, args[1], sig.args[1], types.intp)
|
|
return _np_stack_common(context, builder, sig, args, axis)
|
|
|
|
return sig, codegen
|
|
|
|
|
|
@overload(np.stack)
|
|
def impl_np_stack(arrays, axis=0):
|
|
if isinstance(arrays, types.BaseTuple):
|
|
def impl(arrays, axis=0):
|
|
return np_stack_common(arrays, axis)
|
|
return impl
|
|
|
|
|
|
def NdStack_typer(typingctx, func_name, arrays, ndim_min):
|
|
# does type checking
|
|
dtype, ndim = _sequence_of_arrays(typingctx, func_name, arrays)
|
|
ndim = max(ndim, ndim_min)
|
|
layout = _choose_concatenation_layout(arrays)
|
|
ret = types.Array(dtype, ndim, layout)
|
|
return ret
|
|
|
|
|
|
@intrinsic
|
|
def _np_hstack(typingctx, tup):
|
|
ret = NdStack_typer(typingctx, "np.hstack", tup, 1)
|
|
sig = ret(tup)
|
|
|
|
def codegen(context, builder, sig, args):
|
|
tupty = sig.args[0]
|
|
ndim = tupty[0].ndim
|
|
|
|
if ndim == 0:
|
|
# hstack() on 0-d arrays returns a 1-d array
|
|
axis = context.get_constant(types.intp, 0)
|
|
return _np_stack_common(context, builder, sig, args, axis)
|
|
|
|
else:
|
|
# As a special case, dimension 0 of 1-dimensional arrays
|
|
# is "horizontal"
|
|
axis = 0 if ndim == 1 else 1
|
|
|
|
def np_hstack_impl(arrays):
|
|
return np.concatenate(arrays, axis=axis)
|
|
|
|
return context.compile_internal(builder, np_hstack_impl, sig, args)
|
|
|
|
return sig, codegen
|
|
|
|
|
|
@overload(np.hstack)
|
|
def impl_np_hstack(tup):
|
|
if isinstance(tup, types.BaseTuple):
|
|
def impl(tup):
|
|
return _np_hstack(tup)
|
|
return impl
|
|
|
|
|
|
@intrinsic
|
|
def _np_vstack(typingctx, tup):
|
|
ret = NdStack_typer(typingctx, "np.vstack", tup, 2)
|
|
sig = ret(tup)
|
|
|
|
def codegen(context, builder, sig, args):
|
|
tupty = sig.args[0]
|
|
ndim = tupty[0].ndim
|
|
|
|
if ndim == 0:
|
|
def np_vstack_impl(arrays):
|
|
return np.expand_dims(np.hstack(arrays), 1)
|
|
|
|
elif ndim == 1:
|
|
# np.stack(arrays, axis=0)
|
|
axis = context.get_constant(types.intp, 0)
|
|
return _np_stack_common(context, builder, sig, args, axis)
|
|
|
|
else:
|
|
def np_vstack_impl(arrays):
|
|
return np.concatenate(arrays, axis=0)
|
|
|
|
return context.compile_internal(builder, np_vstack_impl, sig, args)
|
|
|
|
return sig, codegen
|
|
|
|
|
|
@overload(np.vstack)
|
|
def impl_np_vstack(tup):
|
|
if isinstance(tup, types.BaseTuple):
|
|
def impl(tup):
|
|
return _np_vstack(tup)
|
|
return impl
|
|
|
|
|
|
@intrinsic
|
|
def _np_dstack(typingctx, tup):
|
|
ret = NdStack_typer(typingctx, "np.dstack", tup, 3)
|
|
sig = ret(tup)
|
|
|
|
def codegen(context, builder, sig, args):
|
|
tupty = sig.args[0]
|
|
retty = sig.return_type
|
|
ndim = tupty[0].ndim
|
|
|
|
if ndim == 0:
|
|
def np_vstack_impl(arrays):
|
|
return np.hstack(arrays).reshape(1, 1, -1)
|
|
|
|
return context.compile_internal(builder, np_vstack_impl, sig, args)
|
|
|
|
elif ndim == 1:
|
|
# np.expand_dims(np.stack(arrays, axis=1), axis=0)
|
|
axis = context.get_constant(types.intp, 1)
|
|
stack_retty = retty.copy(ndim=retty.ndim - 1)
|
|
stack_sig = typing.signature(stack_retty, *sig.args)
|
|
stack_ret = _np_stack_common(context, builder, stack_sig, args,
|
|
axis)
|
|
|
|
axis = context.get_constant(types.intp, 0)
|
|
expand_sig = typing.signature(retty, stack_retty)
|
|
return expand_dims(context, builder, expand_sig, (stack_ret,), axis)
|
|
|
|
elif ndim == 2:
|
|
# np.stack(arrays, axis=2)
|
|
axis = context.get_constant(types.intp, 2)
|
|
return _np_stack_common(context, builder, sig, args, axis)
|
|
|
|
else:
|
|
def np_vstack_impl(arrays):
|
|
return np.concatenate(arrays, axis=2)
|
|
|
|
return context.compile_internal(builder, np_vstack_impl, sig, args)
|
|
|
|
return sig, codegen
|
|
|
|
|
|
@overload(np.dstack)
|
|
def impl_np_dstack(tup):
|
|
if isinstance(tup, types.BaseTuple):
|
|
def impl(tup):
|
|
return _np_dstack(tup)
|
|
return impl
|
|
|
|
|
|
@extending.overload_method(types.Array, 'fill')
|
|
def arr_fill(arr, val):
|
|
|
|
def fill_impl(arr, val):
|
|
arr[:] = val
|
|
return None
|
|
|
|
return fill_impl
|
|
|
|
|
|
@extending.overload_method(types.Array, 'dot')
|
|
def array_dot(arr, other):
|
|
def dot_impl(arr, other):
|
|
return np.dot(arr, other)
|
|
|
|
return dot_impl
|
|
|
|
|
|
@overload(np.fliplr)
|
|
def np_flip_lr(m):
|
|
|
|
if not type_can_asarray(m):
|
|
raise errors.TypingError("Cannot np.fliplr on %s type" % m)
|
|
|
|
def impl(m):
|
|
A = np.asarray(m)
|
|
# this handling is superfluous/dead as < 2d array cannot be indexed as
|
|
# present below and so typing fails. If the typing doesn't fail due to
|
|
# some future change, this will catch it.
|
|
if A.ndim < 2:
|
|
raise ValueError('Input must be >= 2-d.')
|
|
return A[::, ::-1, ...]
|
|
return impl
|
|
|
|
|
|
@overload(np.flipud)
|
|
def np_flip_ud(m):
|
|
|
|
if not type_can_asarray(m):
|
|
raise errors.TypingError("Cannot np.flipud on %s type" % m)
|
|
|
|
def impl(m):
|
|
A = np.asarray(m)
|
|
# this handling is superfluous/dead as a 0d array cannot be indexed as
|
|
# present below and so typing fails. If the typing doesn't fail due to
|
|
# some future change, this will catch it.
|
|
if A.ndim < 1:
|
|
raise ValueError('Input must be >= 1-d.')
|
|
return A[::-1, ...]
|
|
return impl
|
|
|
|
|
|
@intrinsic
|
|
def _build_flip_slice_tuple(tyctx, sz):
|
|
""" Creates a tuple of slices for np.flip indexing like
|
|
`(slice(None, None, -1),) * sz` """
|
|
if not isinstance(sz, types.IntegerLiteral):
|
|
raise errors.RequireLiteralValue(sz)
|
|
size = int(sz.literal_value)
|
|
tuple_type = types.UniTuple(dtype=types.slice3_type, count=size)
|
|
sig = tuple_type(sz)
|
|
|
|
def codegen(context, builder, signature, args):
|
|
def impl(length, empty_tuple):
|
|
out = empty_tuple
|
|
for i in range(length):
|
|
out = tuple_setitem(out, i, slice(None, None, -1))
|
|
return out
|
|
|
|
inner_argtypes = [types.intp, tuple_type]
|
|
inner_sig = typing.signature(tuple_type, *inner_argtypes)
|
|
ll_idx_type = context.get_value_type(types.intp)
|
|
# Allocate an empty tuple
|
|
empty_tuple = context.get_constant_undef(tuple_type)
|
|
inner_args = [ll_idx_type(size), empty_tuple]
|
|
|
|
res = context.compile_internal(builder, impl, inner_sig, inner_args)
|
|
return res
|
|
|
|
return sig, codegen
|
|
|
|
|
|
@overload(np.flip)
|
|
def np_flip(m):
|
|
# a constant value is needed for the tuple slice, types.Array.ndim can
|
|
# provide this and so at presnet only type.Array is support
|
|
if not isinstance(m, types.Array):
|
|
raise errors.TypingError("Cannot np.flip on %s type" % m)
|
|
|
|
def impl(m):
|
|
sl = _build_flip_slice_tuple(m.ndim)
|
|
return m[sl]
|
|
|
|
return impl
|
|
|
|
|
|
@overload(np.array_split)
|
|
def np_array_split(ary, indices_or_sections, axis=0):
|
|
if isinstance(ary, (types.UniTuple, types.ListType, types.List)):
|
|
def impl(ary, indices_or_sections, axis=0):
|
|
return np.array_split(
|
|
np.asarray(ary),
|
|
indices_or_sections,
|
|
axis=axis
|
|
)
|
|
|
|
return impl
|
|
|
|
if isinstance(indices_or_sections, types.Integer):
|
|
def impl(ary, indices_or_sections, axis=0):
|
|
l, rem = divmod(ary.shape[axis], indices_or_sections)
|
|
indices = np.cumsum(np.array(
|
|
[l + 1] * rem +
|
|
[l] * (indices_or_sections - rem - 1)
|
|
))
|
|
return np.array_split(ary, indices, axis=axis)
|
|
|
|
return impl
|
|
|
|
elif (
|
|
isinstance(indices_or_sections, types.IterableType)
|
|
and isinstance(
|
|
indices_or_sections.iterator_type.yield_type,
|
|
types.Integer
|
|
)
|
|
):
|
|
def impl(ary, indices_or_sections, axis=0):
|
|
slice_tup = build_full_slice_tuple(ary.ndim)
|
|
axis = normalize_axis("np.split", "axis", ary.ndim, axis)
|
|
out = []
|
|
prev = 0
|
|
for cur in indices_or_sections:
|
|
idx = tuple_setitem(slice_tup, axis, slice(prev, cur))
|
|
out.append(ary[idx])
|
|
prev = cur
|
|
out.append(ary[tuple_setitem(slice_tup, axis, slice(cur, None))])
|
|
return out
|
|
|
|
return impl
|
|
|
|
elif (
|
|
isinstance(indices_or_sections, types.Tuple)
|
|
and all(isinstance(t, types.Integer) for t in indices_or_sections.types)
|
|
):
|
|
def impl(ary, indices_or_sections, axis=0):
|
|
slice_tup = build_full_slice_tuple(ary.ndim)
|
|
axis = normalize_axis("np.split", "axis", ary.ndim, axis)
|
|
out = []
|
|
prev = 0
|
|
for cur in literal_unroll(indices_or_sections):
|
|
idx = tuple_setitem(slice_tup, axis, slice(prev, cur))
|
|
out.append(ary[idx])
|
|
prev = cur
|
|
out.append(ary[tuple_setitem(slice_tup, axis, slice(cur, None))])
|
|
return out
|
|
|
|
return impl
|
|
|
|
|
|
@overload(np.split)
|
|
def np_split(ary, indices_or_sections, axis=0):
|
|
# This is just a wrapper of array_split, but with an extra error if
|
|
# indices is an int.
|
|
if isinstance(ary, (types.UniTuple, types.ListType, types.List)):
|
|
def impl(ary, indices_or_sections, axis=0):
|
|
return np.split(np.asarray(ary), indices_or_sections, axis=axis)
|
|
|
|
return impl
|
|
|
|
if isinstance(indices_or_sections, types.Integer):
|
|
def impl(ary, indices_or_sections, axis=0):
|
|
_, rem = divmod(ary.shape[axis], indices_or_sections)
|
|
if rem != 0:
|
|
raise ValueError(
|
|
"array split does not result in an equal division"
|
|
)
|
|
return np.array_split(
|
|
ary, indices_or_sections, axis=axis
|
|
)
|
|
|
|
return impl
|
|
|
|
else:
|
|
return np_array_split(ary, indices_or_sections, axis=axis)
|
|
|
|
|
|
@overload(np.vsplit)
|
|
def numpy_vsplit(ary, indices_or_sections):
|
|
if not isinstance(ary, types.Array):
|
|
msg = 'The argument "ary" must be an array'
|
|
raise errors.TypingError(msg)
|
|
|
|
if not isinstance(indices_or_sections, (types.Integer, types.Array,
|
|
types.List, types.UniTuple)):
|
|
msg = ('The argument "indices_or_sections" must be int or 1d-array')
|
|
raise errors.TypingError(msg)
|
|
|
|
def impl(ary, indices_or_sections):
|
|
if ary.ndim < 2:
|
|
raise ValueError(('vsplit only works on '
|
|
'arrays of 2 or more dimensions'))
|
|
return np.split(ary, indices_or_sections, axis=0)
|
|
|
|
return impl
|
|
|
|
|
|
@overload(np.hsplit)
|
|
def numpy_hsplit(ary, indices_or_sections):
|
|
if not isinstance(ary, types.Array):
|
|
msg = 'The argument "ary" must be an array'
|
|
raise errors.TypingError(msg)
|
|
|
|
if not isinstance(indices_or_sections, (types.Integer, types.Array,
|
|
types.List, types.UniTuple)):
|
|
msg = ('The argument "indices_or_sections" must be int or 1d-array')
|
|
raise errors.TypingError(msg)
|
|
|
|
def impl(ary, indices_or_sections):
|
|
if ary.ndim == 0:
|
|
raise ValueError(('hsplit only works on '
|
|
'arrays of 1 or more dimensions'))
|
|
if ary.ndim > 1:
|
|
return np.split(ary, indices_or_sections, axis=1)
|
|
return np.split(ary, indices_or_sections, axis=0)
|
|
|
|
return impl
|
|
|
|
|
|
@overload(np.dsplit)
|
|
def numpy_dsplit(ary, indices_or_sections):
|
|
if not isinstance(ary, types.Array):
|
|
msg = 'The argument "ary" must be an array'
|
|
raise errors.TypingError(msg)
|
|
|
|
if not isinstance(indices_or_sections, (types.Integer, types.Array,
|
|
types.List, types.UniTuple)):
|
|
msg = ('The argument "indices_or_sections" must be int or 1d-array')
|
|
raise errors.TypingError(msg)
|
|
|
|
def impl(ary, indices_or_sections):
|
|
if ary.ndim < 3:
|
|
raise ValueError('dsplit only works on arrays of 3 or more '
|
|
'dimensions')
|
|
return np.split(ary, indices_or_sections, axis=2)
|
|
|
|
return impl
|
|
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Sorting
|
|
|
|
_sorts = {}
|
|
|
|
|
|
def default_lt(a, b):
|
|
"""
|
|
Trivial comparison function between two keys.
|
|
"""
|
|
return a < b
|
|
|
|
|
|
def get_sort_func(kind, lt_impl, is_argsort=False):
|
|
"""
|
|
Get a sort implementation of the given kind.
|
|
"""
|
|
key = kind, lt_impl.__name__, is_argsort
|
|
|
|
try:
|
|
return _sorts[key]
|
|
except KeyError:
|
|
if kind == 'quicksort':
|
|
sort = quicksort.make_jit_quicksort(
|
|
lt=lt_impl,
|
|
is_argsort=is_argsort,
|
|
is_np_array=True)
|
|
func = sort.run_quicksort
|
|
elif kind == 'mergesort':
|
|
sort = mergesort.make_jit_mergesort(
|
|
lt=lt_impl,
|
|
is_argsort=is_argsort)
|
|
func = sort.run_mergesort
|
|
_sorts[key] = func
|
|
return func
|
|
|
|
|
|
def lt_implementation(dtype):
|
|
if isinstance(dtype, types.Float):
|
|
return lt_floats
|
|
elif isinstance(dtype, types.Complex):
|
|
return lt_complex
|
|
else:
|
|
return default_lt
|
|
|
|
|
|
@lower_builtin("array.sort", types.Array)
|
|
def array_sort(context, builder, sig, args):
|
|
arytype = sig.args[0]
|
|
|
|
sort_func = get_sort_func(kind='quicksort',
|
|
lt_impl=lt_implementation(arytype.dtype))
|
|
|
|
def array_sort_impl(arr):
|
|
# Note we clobber the return value
|
|
sort_func(arr)
|
|
|
|
return context.compile_internal(builder, array_sort_impl, sig, args)
|
|
|
|
|
|
@overload(np.sort)
|
|
def impl_np_sort(a):
|
|
if not type_can_asarray(a):
|
|
raise errors.TypingError('Argument "a" must '
|
|
'be array-like')
|
|
|
|
def np_sort_impl(a):
|
|
res = a.copy()
|
|
res.sort()
|
|
return res
|
|
return np_sort_impl
|
|
|
|
|
|
@lower_builtin("array.argsort", types.Array, types.StringLiteral)
|
|
@lower_builtin(np.argsort, types.Array, types.StringLiteral)
|
|
def array_argsort(context, builder, sig, args):
|
|
arytype, kind = sig.args
|
|
|
|
sort_func = get_sort_func(kind=kind.literal_value,
|
|
lt_impl=lt_implementation(arytype.dtype),
|
|
is_argsort=True)
|
|
|
|
def array_argsort_impl(arr):
|
|
return sort_func(arr)
|
|
|
|
innersig = sig.replace(args=sig.args[:1])
|
|
innerargs = args[:1]
|
|
return context.compile_internal(builder, array_argsort_impl,
|
|
innersig, innerargs)
|
|
|
|
|
|
# ------------------------------------------------------------------------------
|
|
# Implicit cast
|
|
|
|
@lower_cast(types.Array, types.Array)
|
|
def array_to_array(context, builder, fromty, toty, val):
|
|
# Type inference should have prevented illegal array casting.
|
|
assert fromty.mutable != toty.mutable or toty.layout == 'A'
|
|
return val
|
|
|
|
|
|
@lower_cast(types.Array, types.UnicodeCharSeq)
|
|
@lower_cast(types.Array, types.Float)
|
|
@lower_cast(types.Array, types.Integer)
|
|
@lower_cast(types.Array, types.Complex)
|
|
@lower_cast(types.Array, types.Boolean)
|
|
@lower_cast(types.Array, types.NPTimedelta)
|
|
@lower_cast(types.Array, types.NPDatetime)
|
|
def array0d_to_scalar(context, builder, fromty, toty, val):
|
|
def impl(a):
|
|
# a is an array(T, 0d, O), T is type, O is order
|
|
return a.take(0)
|
|
|
|
sig = signature(toty, fromty)
|
|
res = context.compile_internal(builder, impl, sig, [val])
|
|
return impl_ret_untracked(context, builder, sig.return_type, res)
|
|
|
|
|
|
@lower_cast(types.Array, types.UnicodeCharSeq)
|
|
def array_to_unichrseq(context, builder, fromty, toty, val):
|
|
def impl(a):
|
|
return str(a[()])
|
|
|
|
sig = signature(toty, fromty)
|
|
res = context.compile_internal(builder, impl, sig, [val])
|
|
return impl_ret_borrowed(context, builder, sig.return_type, res)
|
|
|
|
|
|
# ------------------------------------------------------------------------------
|
|
# Stride tricks
|
|
|
|
def reshape_unchecked(a, shape, strides):
|
|
"""
|
|
An intrinsic returning a derived array with the given shape and strides.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
|
|
@extending.type_callable(reshape_unchecked)
|
|
def type_reshape_unchecked(context):
|
|
def check_shape(shape):
|
|
return (isinstance(shape, types.BaseTuple) and
|
|
all(isinstance(v, types.Integer) for v in shape))
|
|
|
|
def typer(a, shape, strides):
|
|
if not isinstance(a, types.Array):
|
|
return
|
|
if not check_shape(shape) or not check_shape(strides):
|
|
return
|
|
if len(shape) != len(strides):
|
|
return
|
|
return a.copy(ndim=len(shape), layout='A')
|
|
|
|
return typer
|
|
|
|
|
|
@lower_builtin(reshape_unchecked, types.Array, types.BaseTuple, types.BaseTuple)
|
|
def impl_shape_unchecked(context, builder, sig, args):
|
|
aryty = sig.args[0]
|
|
retty = sig.return_type
|
|
|
|
ary = make_array(aryty)(context, builder, args[0])
|
|
out = make_array(retty)(context, builder)
|
|
shape = cgutils.unpack_tuple(builder, args[1])
|
|
strides = cgutils.unpack_tuple(builder, args[2])
|
|
|
|
populate_array(out,
|
|
data=ary.data,
|
|
shape=shape,
|
|
strides=strides,
|
|
itemsize=ary.itemsize,
|
|
meminfo=ary.meminfo,
|
|
)
|
|
|
|
res = out._getvalue()
|
|
return impl_ret_borrowed(context, builder, retty, res)
|
|
|
|
|
|
@extending.overload(np.lib.stride_tricks.as_strided)
|
|
def as_strided(x, shape=None, strides=None):
|
|
if shape in (None, types.none):
|
|
@register_jitable
|
|
def get_shape(x, shape):
|
|
return x.shape
|
|
else:
|
|
@register_jitable
|
|
def get_shape(x, shape):
|
|
return shape
|
|
|
|
if strides in (None, types.none):
|
|
# When *strides* is not passed, as_strided() does a non-size-checking
|
|
# reshape(), possibly changing the original strides. This is too
|
|
# cumbersome to support right now, and a Web search shows all example
|
|
# use cases of as_strided() pass explicit *strides*.
|
|
raise NotImplementedError("as_strided() strides argument is mandatory")
|
|
else:
|
|
@register_jitable
|
|
def get_strides(x, strides):
|
|
return strides
|
|
|
|
def as_strided_impl(x, shape=None, strides=None):
|
|
x = reshape_unchecked(x, get_shape(x, shape), get_strides(x, strides))
|
|
return x
|
|
|
|
return as_strided_impl
|
|
|
|
|
|
@extending.overload(np.lib.stride_tricks.sliding_window_view)
|
|
def sliding_window_view(x, window_shape, axis=None):
|
|
|
|
# Window shape must be given as either an integer or tuple of integers.
|
|
# We also need to generate buffer tuples we can modify to contain the
|
|
# final shape and strides (reshape_unchecked does not accept lists).
|
|
if isinstance(window_shape, types.Integer):
|
|
shape_buffer = tuple(range(x.ndim + 1))
|
|
stride_buffer = tuple(range(x.ndim + 1))
|
|
|
|
@register_jitable
|
|
def get_window_shape(window_shape):
|
|
return (window_shape,)
|
|
|
|
elif (isinstance(window_shape, types.UniTuple) and
|
|
isinstance(window_shape.dtype, types.Integer)):
|
|
shape_buffer = tuple(range(x.ndim + len(window_shape)))
|
|
stride_buffer = tuple(range(x.ndim + len(window_shape)))
|
|
|
|
@register_jitable
|
|
def get_window_shape(window_shape):
|
|
return window_shape
|
|
|
|
else:
|
|
raise errors.TypingError(
|
|
"window_shape must be an integer or tuple of integers"
|
|
)
|
|
|
|
# Axis must be integer, tuple of integers, or None for all axes.
|
|
if is_nonelike(axis):
|
|
@register_jitable
|
|
def get_axis(window_shape, axis, ndim):
|
|
return list(range(ndim))
|
|
|
|
elif isinstance(axis, types.Integer):
|
|
@register_jitable
|
|
def get_axis(window_shape, axis, ndim):
|
|
return [
|
|
normalize_axis("sliding_window_view", "axis", ndim, axis)
|
|
]
|
|
|
|
elif (isinstance(axis, types.UniTuple) and
|
|
isinstance(axis.dtype, types.Integer)):
|
|
@register_jitable
|
|
def get_axis(window_shape, axis, ndim):
|
|
return [normalize_axis("sliding_window_view", "axis", ndim, a)
|
|
for a in axis]
|
|
|
|
else:
|
|
raise errors.TypingError(
|
|
"axis must be None, an integer or tuple of integers"
|
|
)
|
|
|
|
def sliding_window_view_impl(x, window_shape, axis=None):
|
|
window_shape = get_window_shape(window_shape)
|
|
axis = get_axis(window_shape, axis, x.ndim)
|
|
if len(window_shape) != len(axis):
|
|
raise ValueError(
|
|
"Must provide matching length window_shape and axis"
|
|
)
|
|
|
|
# Initialise view details with shape and strides of x.
|
|
out_shape = shape_buffer
|
|
out_strides = stride_buffer
|
|
for i in range(x.ndim):
|
|
out_shape = tuple_setitem(out_shape, i, x.shape[i])
|
|
out_strides = tuple_setitem(out_strides, i, x.strides[i])
|
|
|
|
# Trim the dimensions being windowed and set the window shape and
|
|
# strides. Note: the same axis can be windowed repeatedly.
|
|
i = x.ndim
|
|
for ax, dim in zip(axis, window_shape):
|
|
if dim < 0:
|
|
raise ValueError(
|
|
"`window_shape` cannot contain negative values"
|
|
)
|
|
if out_shape[ax] < dim:
|
|
raise ValueError(
|
|
"window_shape cannot be larger than input array shape"
|
|
)
|
|
|
|
trimmed = out_shape[ax] - dim + 1
|
|
out_shape = tuple_setitem(out_shape, ax, trimmed)
|
|
out_shape = tuple_setitem(out_shape, i, dim)
|
|
out_strides = tuple_setitem(out_strides, i, x.strides[ax])
|
|
i += 1
|
|
|
|
# The NumPy version calls as_strided, but our implementation of
|
|
# as_strided is effectively a wrapper for reshape_unchecked.
|
|
view = reshape_unchecked(x, out_shape, out_strides)
|
|
return view
|
|
|
|
return sliding_window_view_impl
|
|
|
|
|
|
@overload(bool)
|
|
def ol_bool(arr):
|
|
if isinstance(arr, types.Array):
|
|
def impl(arr):
|
|
if arr.size == 0:
|
|
return False # this is deprecated
|
|
elif arr.size == 1:
|
|
return bool(arr.take(0))
|
|
else:
|
|
msg = ("The truth value of an array with more than one element "
|
|
"is ambiguous. Use a.any() or a.all()")
|
|
raise ValueError(msg)
|
|
return impl
|
|
|
|
|
|
@overload(np.swapaxes)
|
|
def numpy_swapaxes(a, axis1, axis2):
|
|
if not isinstance(axis1, (int, types.Integer)):
|
|
raise errors.TypingError('The second argument "axis1" must be an '
|
|
'integer')
|
|
if not isinstance(axis2, (int, types.Integer)):
|
|
raise errors.TypingError('The third argument "axis2" must be an '
|
|
'integer')
|
|
if not isinstance(a, types.Array):
|
|
raise errors.TypingError('The first argument "a" must be an array')
|
|
|
|
# create tuple list for transpose
|
|
ndim = a.ndim
|
|
axes_list = tuple(range(ndim))
|
|
|
|
def impl(a, axis1, axis2):
|
|
axis1 = normalize_axis("np.swapaxes", "axis1", ndim, axis1)
|
|
axis2 = normalize_axis("np.swapaxes", "axis2", ndim, axis2)
|
|
|
|
# to ensure tuple_setitem support of negative values
|
|
if axis1 < 0:
|
|
axis1 += ndim
|
|
if axis2 < 0:
|
|
axis2 += ndim
|
|
|
|
axes_tuple = tuple_setitem(axes_list, axis1, axis2)
|
|
axes_tuple = tuple_setitem(axes_tuple, axis2, axis1)
|
|
return np.transpose(a, axes_tuple)
|
|
|
|
return impl
|
|
|
|
|
|
@register_jitable
|
|
def _take_along_axis_impl(
|
|
arr, indices, axis, Ni_orig, Nk_orig, indices_broadcast_shape
|
|
):
|
|
# Based on example code in
|
|
# https://github.com/numpy/numpy/blob/623bc1fae1d47df24e7f1e29321d0c0ba2771ce0/numpy/lib/shape_base.py#L90-L103
|
|
# With addition of pre-broadcasting:
|
|
# https://github.com/numpy/numpy/issues/19704
|
|
|
|
# Wrap axis, it's used in tuple_setitem so must be (axis >= 0) to ensure
|
|
# the GEP is in bounds.
|
|
axis = normalize_axis("np.take_along_axis", "axis", arr.ndim, axis)
|
|
|
|
# Broadcast the two arrays to matching shapes:
|
|
arr_shape = list(arr.shape)
|
|
arr_shape[axis] = 1
|
|
for i, (d1, d2) in enumerate(zip(arr_shape, indices.shape)):
|
|
if d1 == 1:
|
|
new_val = d2
|
|
elif d2 == 1:
|
|
new_val = d1
|
|
else:
|
|
if d1 != d2:
|
|
raise ValueError(
|
|
"`arr` and `indices` dimensions don't match"
|
|
)
|
|
new_val = d1
|
|
indices_broadcast_shape = tuple_setitem(
|
|
indices_broadcast_shape, i, new_val
|
|
)
|
|
arr_broadcast_shape = tuple_setitem(
|
|
indices_broadcast_shape, axis, arr.shape[axis]
|
|
)
|
|
arr = np.broadcast_to(arr, arr_broadcast_shape)
|
|
indices = np.broadcast_to(indices, indices_broadcast_shape)
|
|
|
|
Ni = Ni_orig
|
|
if len(Ni_orig) > 0:
|
|
for i in range(len(Ni)):
|
|
Ni = tuple_setitem(Ni, i, arr.shape[i])
|
|
Nk = Nk_orig
|
|
if len(Nk_orig) > 0:
|
|
for i in range(len(Nk)):
|
|
Nk = tuple_setitem(Nk, i, arr.shape[axis + 1 + i])
|
|
|
|
J = indices.shape[axis] # Need not equal M
|
|
out = np.empty(Ni + (J,) + Nk, arr.dtype)
|
|
|
|
np_s_ = (slice(None, None, None),)
|
|
|
|
for ii in np.ndindex(Ni):
|
|
for kk in np.ndindex(Nk):
|
|
a_1d = arr[ii + np_s_ + kk]
|
|
indices_1d = indices[ii + np_s_ + kk]
|
|
out_1d = out[ii + np_s_ + kk]
|
|
for j in range(J):
|
|
out_1d[j] = a_1d[indices_1d[j]]
|
|
return out
|
|
|
|
|
|
@overload(np.take_along_axis)
|
|
def arr_take_along_axis(arr, indices, axis):
|
|
if not isinstance(arr, types.Array):
|
|
raise errors.TypingError('The first argument "arr" must be an array')
|
|
if not isinstance(indices, types.Array):
|
|
raise errors.TypingError(
|
|
'The second argument "indices" must be an array')
|
|
if not isinstance(indices.dtype, types.Integer):
|
|
raise errors.TypingError('The indices array must contain integers')
|
|
if is_nonelike(axis):
|
|
arr_ndim = 1
|
|
else:
|
|
arr_ndim = arr.ndim
|
|
if arr_ndim != indices.ndim:
|
|
# Matches NumPy error:
|
|
raise errors.TypingError(
|
|
"`indices` and `arr` must have the same number of dimensions"
|
|
)
|
|
|
|
indices_broadcast_shape = tuple(range(indices.ndim))
|
|
if is_nonelike(axis):
|
|
def take_along_axis_impl(arr, indices, axis):
|
|
return _take_along_axis_impl(arr.flatten(), indices, 0, (), (),
|
|
indices_broadcast_shape)
|
|
else:
|
|
check_is_integer(axis, "axis")
|
|
if not isinstance(axis, types.IntegerLiteral):
|
|
raise errors.NumbaValueError("axis must be a literal value")
|
|
axis = axis.literal_value
|
|
if axis < 0:
|
|
axis = arr.ndim + axis
|
|
|
|
if axis < 0 or axis >= arr.ndim:
|
|
raise errors.NumbaValueError("axis is out of bounds")
|
|
|
|
Ni = tuple(range(axis))
|
|
Nk = tuple(range(axis + 1, arr.ndim))
|
|
|
|
def take_along_axis_impl(arr, indices, axis):
|
|
return _take_along_axis_impl(arr, indices, axis, Ni, Nk,
|
|
indices_broadcast_shape)
|
|
return take_along_axis_impl
|
|
|
|
|
|
@overload(np.nan_to_num)
|
|
def nan_to_num_impl(x, copy=True, nan=0.0):
|
|
if isinstance(x, types.Number):
|
|
if isinstance(x, types.Integer):
|
|
# Integers do not have nans or infs
|
|
def impl(x, copy=True, nan=0.0):
|
|
return x
|
|
|
|
elif isinstance(x, types.Float):
|
|
def impl(x, copy=True, nan=0.0):
|
|
if np.isnan(x):
|
|
return nan
|
|
elif np.isneginf(x):
|
|
return np.finfo(type(x)).min
|
|
elif np.isposinf(x):
|
|
return np.finfo(type(x)).max
|
|
return x
|
|
elif isinstance(x, types.Complex):
|
|
def impl(x, copy=True, nan=0.0):
|
|
r = np.nan_to_num(x.real, nan=nan)
|
|
c = np.nan_to_num(x.imag, nan=nan)
|
|
return complex(r, c)
|
|
else:
|
|
raise errors.TypingError(
|
|
"Only Integer, Float, and Complex values are accepted"
|
|
)
|
|
|
|
elif type_can_asarray(x):
|
|
if isinstance(x.dtype, types.Integer):
|
|
# Integers do not have nans or infs
|
|
def impl(x, copy=True, nan=0.0):
|
|
return x
|
|
elif isinstance(x.dtype, types.Float):
|
|
def impl(x, copy=True, nan=0.0):
|
|
min_inf = np.finfo(x.dtype).min
|
|
max_inf = np.finfo(x.dtype).max
|
|
|
|
x_ = np.asarray(x)
|
|
output = np.copy(x_) if copy else x_
|
|
|
|
output_flat = output.flat
|
|
for i in range(output.size):
|
|
if np.isnan(output_flat[i]):
|
|
output_flat[i] = nan
|
|
elif np.isneginf(output_flat[i]):
|
|
output_flat[i] = min_inf
|
|
elif np.isposinf(output_flat[i]):
|
|
output_flat[i] = max_inf
|
|
return output
|
|
elif isinstance(x.dtype, types.Complex):
|
|
def impl(x, copy=True, nan=0.0):
|
|
x_ = np.asarray(x)
|
|
output = np.copy(x_) if copy else x_
|
|
|
|
np.nan_to_num(output.real, copy=False, nan=nan)
|
|
np.nan_to_num(output.imag, copy=False, nan=nan)
|
|
return output
|
|
else:
|
|
raise errors.TypingError(
|
|
"Only Integer, Float, and Complex values are accepted"
|
|
)
|
|
else:
|
|
raise errors.TypingError("The first argument must be a scalar or an "
|
|
"array-like")
|
|
return impl
|