numba_vulkan.arrayfuncs

Reductions over arrays, written in Python on top of element access.

Numba’s own sum, min… walk the array through its data pointer, which arrays on this target do not have. The versions here are overloads for the vulkan target that read the elements by position instead (see numba_vulkan.vkimpl.flat_item), so they work for arrays and views of any dimensionality. They reduce over all elements, or, with a constant axis, along one axis: kernels cannot allocate the resulting array, so a.sum(axis=0) is an array expression (see AxisReduction), whose elements are computed where they are read.

Each is available as a method (a.sum()) and as the NumPy function (np.sum(a)).

Attributes

TARGET

Classes

AxisReduction

The operation of an array expression that reduces along an axis.

Module Contents

numba_vulkan.arrayfuncs.TARGET = 'vulkan'

class numba_vulkan.arrayfuncs.AxisReduction

The operation of an array expression that reduces along an axis.

a.sum(axis=k) is a numba_vulkan.vktypes.VulkanExpr with this as its operation, and a and the axis as its operands; the axis is a run-time value, as in Numba’s sum(axis=...) on the CPU, since the number of dimensions of the result does not depend on it. Its shape is that of a without the axis; each element is a loop along the axis, written in Python and compiled where the element is read (see numba_vulkan.vkimpl.expression_element).

Parameters:

• name (str) – The reduction: a key of _REDUCTIONS.

• acc (numba.types.Type) – Type in which sums, products and means are formed.

name: str

acc: numba.core.types.Type

reduces_axis = True

element_function(ndim)

The function that computes one element.

Parameters:

ndim (int) – Dimensions of the operand.

Returns:

Takes the operand, the axis (made non-negative, and checked when the shape was computed) and one index per dimension of the result.

Return type:

function