import inspect
from numba.np.ufunc import _internal
from numba.np.ufunc.parallel import ParallelUFuncBuilder, ParallelGUFuncBuilder
from numba.core.registry import DelayedRegistry
from numba.np.ufunc import dufunc
from numba.np.ufunc import gufunc
class _BaseVectorize(object):
@classmethod
def get_identity(cls, kwargs):
return kwargs.pop('identity', None)
@classmethod
def get_cache(cls, kwargs):
return kwargs.pop('cache', False)
@classmethod
def get_writable_args(cls, kwargs):
return kwargs.pop('writable_args', ())
@classmethod
def get_target_implementation(cls, kwargs):
target = kwargs.pop('target', 'cpu')
try:
return cls.target_registry[target]
except KeyError:
raise ValueError("Unsupported target: %s" % target)
class Vectorize(_BaseVectorize):
target_registry = DelayedRegistry({'cpu': dufunc.DUFunc,
'parallel': ParallelUFuncBuilder,})
def __new__(cls, func, **kws):
identity = cls.get_identity(kws)
cache = cls.get_cache(kws)
imp = cls.get_target_implementation(kws)
return imp(func, identity=identity, cache=cache, targetoptions=kws)
class GUVectorize(_BaseVectorize):
target_registry = DelayedRegistry({'cpu': gufunc.GUFunc,
'parallel': ParallelGUFuncBuilder,})
def __new__(cls, func, signature, **kws):
identity = cls.get_identity(kws)
cache = cls.get_cache(kws)
imp = cls.get_target_implementation(kws)
writable_args = cls.get_writable_args(kws)
if imp is gufunc.GUFunc:
is_dyn = kws.pop('is_dynamic', False)
return imp(func, signature, identity=identity, cache=cache,
is_dynamic=is_dyn, targetoptions=kws,
writable_args=writable_args)
else:
return imp(func, signature, identity=identity, cache=cache,
targetoptions=kws, writable_args=writable_args)
[docs]def vectorize(ftylist_or_function=(), **kws):
"""vectorize(ftylist_or_function=(), target='cpu', identity=None, **kws)
A decorator that creates a NumPy ufunc object using Numba compiled
code. When no arguments or only keyword arguments are given,
vectorize will return a Numba dynamic ufunc (DUFunc) object, where
compilation/specialization may occur at call-time.
Args
-----
ftylist_or_function: function or iterable
When the first argument is a function, signatures are dealt
with at call-time.
When the first argument is an iterable of type signatures,
which are either function type object or a string describing
the function type, signatures are finalized at decoration
time.
Keyword Args
------------
target: str
A string for code generation target. Default to "cpu".
identity: int, str, or None
The identity (or unit) value for the element-wise function
being implemented. Allowed values are None (the default), 0, 1,
and "reorderable".
cache: bool
Turns on caching.
Returns
--------
A NumPy universal function
Examples
-------
@vectorize(['float32(float32, float32)',
'float64(float64, float64)'], identity=0)
def sum(a, b):
return a + b
@vectorize
def sum(a, b):
return a + b
@vectorize(identity=1)
def mul(a, b):
return a * b
"""
if isinstance(ftylist_or_function, str):
# Common user mistake
ftylist = [ftylist_or_function]
elif inspect.isfunction(ftylist_or_function):
return dufunc.DUFunc(ftylist_or_function, **kws)
elif ftylist_or_function is not None:
ftylist = ftylist_or_function
def wrap(func):
vec = Vectorize(func, **kws)
for sig in ftylist:
vec.add(sig)
if len(ftylist) > 0:
vec.disable_compile()
return vec.build_ufunc()
return wrap
def guvectorize(*args, **kwargs):
"""guvectorize(ftylist, signature, target='cpu', identity=None, **kws)
A decorator to create NumPy generalized-ufunc object from Numba compiled
code.
Args
-----
ftylist: iterable
An iterable of type signatures, which are either
function type object or a string describing the
function type.
signature: str
A NumPy generalized-ufunc signature.
e.g. "(m, n), (n, p)->(m, p)"
identity: int, str, or None
The identity (or unit) value for the element-wise function
being implemented. Allowed values are None (the default), 0, 1,
and "reorderable".
cache: bool
Turns on caching.
writable_args: tuple
a tuple of indices of input variables that are writable.
target: str
A string for code generation target. Defaults to "cpu".
Returns
--------
A NumPy generalized universal-function
Example
-------
@guvectorize(['void(int32[:,:], int32[:,:], int32[:,:])',
'void(float32[:,:], float32[:,:], float32[:,:])'],
'(x, y),(x, y)->(x, y)')
def add_2d_array(a, b, c):
for i in range(c.shape[0]):
for j in range(c.shape[1]):
c[i, j] = a[i, j] + b[i, j]
"""
if len(args) == 1:
ftylist = []
signature = args[0]
kwargs.setdefault('is_dynamic', True)
elif len(args) == 2:
ftylist = args[0]
signature = args[1]
else:
raise TypeError('guvectorize() takes one or two positional arguments')
if isinstance(ftylist, str):
# Common user mistake
ftylist = [ftylist]
def wrap(func):
guvec = GUVectorize(func, signature, **kwargs)
for fty in ftylist:
guvec.add(fty)
if len(ftylist) > 0:
guvec.disable_compile()
return guvec.build_ufunc()
return wrap