def _create_astype_ufunc(dtype):
name = 'astype_{}'.format(dtype)
rules = tuple([
'{}->{}'.format(cast_from.char, dtype.char)
for cast_from in _dtype_list
])
command = 'out0 = static_cast< {} >(in0)'.format(_dtype_to_ctype[dtype])
return core.create_ufunc(name, rules, command)
def _set_dtype_to_astype_dict():
"""Set a dict with dtypes and astype ufuncs to `_dtype_to_astype_dict`.
Creates a ufunc for type cast operations, and set a dict with keys
as the dtype of the output array and values as astype ufuncs.
This function is called at most once.
"""
global _dtype_to_astype_dict
_dtype_to_astype_dict = {}
dtype_list = [numpy.dtype(type_char) for type_char in '?bhilqBHILQefdFD']
for t in dtype_list:
name = 'astype_{}'.format(t)
rules = tuple(['{}->{}'.format(s.char, t.char) for s in dtype_list])
command = 'out0 = static_cast< {} >(in0)'.format(get_typename(t))
_dtype_to_astype_dict[t] = core.create_ufunc(name, rules, command)
from cupy.core.core import create_ufunc
rsqrt = create_ufunc('cupy_rsqrt', ('e->e', 'f->f', 'd->d', 'F->F', 'D->D'),
'out0 = rsqrt(in0)',
doc='''Returns the reciprocal square root.''')