def impl_keepdims(reduce_op, x, axis, keepdims=False):
axis = np.atleast_1d(np.asarray(axis))
mask = np.zeros(x.ndim, dtype=np.bool8)
mask[axis] = True
original_shape = np.array(x.shape)
squeezed_shape = original_shape[~mask]
# this could be reversed, but we are calling a reduction op on it anyway
new_axes = -np.arange(1, axis.size + 1)
# not that this will copy if reduction happens along a non-contigous axis
x_work = np.moveaxis(x, axis, new_axes)
x_work = np.ascontiguousarray(x_work)
total_reduce = np.prod(original_shape[axis])
total_keep = np.prod(squeezed_shape)
tmp_shape = to_fixed_tuple(np.array((total_keep, total_reduce)), 2)
x_work = np.reshape(x_work, tmp_shape)
result = np.empty((total_keep, ), dtype=x_work.dtype)
for idx in range(result.size):
result[idx] = reduce_op(x_work[idx, ...])
new_shape = original_shape.copy()
new_shape[axis] = 1
new_shape_tuple = to_fixed_tuple(new_shape, x.ndim)
return np.reshape(result, new_shape_tuple)
def unravel_index(index, shape):
sizes = np.zeros(len(shape), dtype=np.int64)
result = np.zeros(len(shape), dtype=np.int64)
sizes[-1] = 1
for i in range(len(shape) - 2, -1, -1):
sizes[i] = sizes[i + 1] * shape[i + 1]
remainder = index
for i in range(len(shape)):
result[i] = remainder // sizes[i]
remainder %= sizes[i]
return to_fixed_tuple(result, len(shape))
def kernel(lhs, rhs):
location = cuda.grid(1)
n_elements = 1
for i in range(lhs.ndim):
n_elements *= lhs.shape[i]
if location >= n_elements:
# bake n_elements into the kernel, better than passing it in
# as another argument.
return
# [0, :] is the to-index (into `lhs`)
# [1, :] is the from-index (into `rhs`)
idx = cuda.local.array(shape=(2, ndim), dtype=types.int64)
for i in range(ndim - 1, -1, -1):
idx[0, i] = location % lhs.shape[i]
idx[1, i] = (location % lhs.shape[i]) * (rhs.shape[i] > 1)
location //= lhs.shape[i]
lhs[to_fixed_tuple(idx[0], ndim)] = rhs[to_fixed_tuple(idx[1], ndim)]
def impl_int(a: np.ndarray, source, destination):
if abs(source) > a.ndim:
raise ValueError("Invalid axis in `source`.")
if abs(destination) > a.ndim:
raise ValueError("Invalid axis in `destination`.")
source = source % a.ndim
destination = destination % a.ndim
order = [n for n in range(a.ndim) if n != source]
order.insert(destination, source)
oder_tuple = to_fixed_tuple(np.array(order), a.ndim)
return np.transpose(a, oder_tuple)
def dirichlet_rv(rng, size, dtype, alphas):
if size_len > 0:
size_tpl = numba_ndarray.to_fixed_tuple(size, size_len)
if (0 < alphas.ndim - 1 <= len(size_tpl)
and size_tpl[neg_ind_shape_len:] != alphas.shape[:-1]):
raise ValueError("Parameters shape and size do not match.")
samples_shape = size_tpl + alphas.shape[-1:]
else:
samples_shape = alphas.shape
res = np.empty(samples_shape, dtype=out_dtype)
alphas_bcast = np.broadcast_to(alphas, samples_shape)
for index in np.ndindex(*samples_shape[:-1]):
res[index] = np.random.dirichlet(alphas_bcast[index])
return (rng, res)
def impl_array(a: np.ndarray, source, destination):
source_work = np.atleast_1d(np.asarray(source))
destination_work = np.atleast_1d(np.asarray(destination))
if source_work.size != destination_work.size:
raise ValueError("`source` and `destination` arguments must have "
"the same number of elements")
for idx in range(source_work.size):
if abs(source_work[idx]) > a.ndim:
raise ValueError("Invalid axis in `source`.")
if abs(destination_work[idx]) > a.ndim:
raise ValueError("Invalid axis in `destination`.")
source_work = [x % a.ndim for x in source_work]
destination_work = [x % a.ndim for x in destination_work]
order = [n for n in range(a.ndim) if n not in source_work]
for dest, src in sorted(zip(destination_work, source_work)):
order.insert(dest, src)
oder_tuple = to_fixed_tuple(np.array(order), a.ndim)
return np.transpose(a, oder_tuple)
def categorical_rv(rng, size, dtype, p):
size_tpl = numba_ndarray.to_fixed_tuple(size, size_len)
ind_shape = p.shape[:-1]
if ind_shape_len > 0:
if size_len > 0 and size_tpl[neg_ind_shape_len:] != ind_shape:
raise ValueError("Parameters shape and size do not match.")
samples_shape = size_tpl[:neg_ind_shape_len] + ind_shape
p_bcast = np.broadcast_to(p,
size_tpl[:neg_ind_shape_len] + p.shape)
else:
samples_shape = size_tpl
p_bcast = p
unif_samples = np.random.uniform(0, 1, samples_shape)
res = np.empty(samples_shape, dtype=out_dtype)
for idx in np.ndindex(*samples_shape):
res[idx] = np.searchsorted(np.cumsum(p_bcast[idx]),
unif_samples[idx])
return (rng, res)
def reshape(x, shape):
new_shape = to_fixed_tuple(shape, ndim)
return np.reshape(x, new_shape)
def tuple_with_length(array, length):
return to_fixed_tuple(array, length)
def foo(array):
a = to_fixed_tuple(array, length=1)
b = to_fixed_tuple(array, 2)
c = to_fixed_tuple(array, const)
d = to_fixed_tuple(array, 0)
return a, b, c, d
def tuple_converter(to_convert):
return to_fixed_tuple(to_convert, tuple_size)
def pair_to_tuple(pair):
return to_fixed_tuple(pair, 2)
def reshape(x, shape):
# TODO: Use this until https://github.com/numba/numba/issues/7353 is closed.
return np.reshape(
np.ascontiguousarray(np.asarray(x)),
numba_ndarray.to_fixed_tuple(shape, ndim),
)
def dirichlet_rv(rng, size, dtype, alphas):
size = numba_ndarray.to_fixed_tuple(size, size_len)
return (rng, np.random.dirichlet(alphas, size))
