def create_math_ufunc(math_name, nargs, name, doc, support_complex=True):
assert 1 <= nargs <= 2
if nargs == 1:
types = ('e->e', 'f->f', 'd->d')
if support_complex:
types += ('F->F', 'D->D')
return _core.create_ufunc(name,
types,
'out0 = %s(in0)' % math_name,
doc=doc)
else:
types = ('ee->e', 'ff->f', 'dd->d')
if support_complex:
types += ('FF->F', 'DD->D')
return _core.create_ufunc(name,
types,
'out0 = %s(in0, in1)' % math_name,
doc=doc)
from cupy import _core
from cupy._core import fusion
add = _core.add
reciprocal = _core.create_ufunc(
'cupy_reciprocal',
('b', 'B', 'h', 'H', 'i', 'I', 'l', 'L', 'q', 'Q', ('e', 'out0 = 1 / in0'),
('f', 'out0 = 1 / in0'), ('d', 'out0 = 1 / in0'),
('F', 'out0 = in0_type(1) / in0'), ('D', 'out0 = in0_type(1) / in0')),
'out0 = in0 == 0 ? 0 : (1 / in0)',
doc='''Computes ``1 / x`` elementwise.
.. seealso:: :data:`numpy.reciprocal`
''')
negative = _core.negative
conjugate = _core.conjugate
angle = _core.angle
def real(val):
'''Returns the real part of the elements of the array.
.. seealso:: :func:`numpy.real`
'''
if fusion._is_fusing():
"""
if fusion._is_fusing():
return fusion._call_ufunc(_math.clip, a, a_min, a_max, out=out)
# TODO(okuta): check type
return a.clip(a_min, a_max, out=out)
# sqrt_fixed is deprecated.
# numpy.sqrt is fixed in numpy 1.11.2.
sqrt = sqrt_fixed = _core.sqrt
cbrt = _core.create_ufunc('cupy_cbrt', ('e->e', 'f->f', 'd->d'),
'out0 = cbrt(in0)',
doc='''Elementwise cube root function.
.. seealso:: :data:`numpy.cbrt`
''')
square = _core.create_ufunc(
'cupy_square',
('b->b', 'B->B', 'h->h', 'H->H', 'i->i', 'I->I', 'l->l', 'L->L', 'q->q',
'Q->Q', 'e->e', 'f->f', 'd->d', 'F->F', 'D->D'),
'out0 = in0 * in0',
doc='''Elementwise square function.
.. seealso:: :data:`numpy.square`
''')
while (x > 2.0) {
x -= 1.0;
y += 1.0 / x;
}
}
if ((1.0 <= x) && (x <= 2.0)) {
y += digamma_imp_1_2(x);
return y;
}
/* x is large, use the asymptotic series */
y += psi_asy(x);
return y;
}
'''
digamma = _core.create_ufunc('cupyx_scipy_digamma', ('f->f', 'd->d'),
'out0 = psi(in0)',
preamble=polevl_definition + psi_definition,
doc="""The digamma function.
Args:
x (cupy.ndarray): The input of digamma function.
Returns:
cupy.ndarray: Computed value of digamma function.
.. seealso:: :data:`scipy.special.digamma`
""")
return z*cevalpoly(coeffs, 22, z);
}
""")
loggamma = _core.create_ufunc(
'cupyx_scipy_loggamma',
(('f->f', 'out0 = out0_type(loggamma_real(in0))'),
('d->d', 'out0 = loggamma_real(in0)'), 'F->F', 'D->D'),
'out0 = out0_type(loggamma(in0))',
preamble=loggamma_definition,
doc="""Principal branch of the logarithm of the gamma function.
Parameters
----------
z : cupy.ndarray
Values in the complex plain at which to compute loggamma
out : cupy.ndarray, optional
Output array for computed values of loggamma
Returns
-------
cupy.ndarray
Values of loggamma at z.
See Also
--------
:func:`scipy.special.loggamma`
""",
)
// pmv_wrap as in
// https://github.com/scipy/scipy/blob/master/scipy/special/specfun_wrappers.c
__device__ double pmv_wrap(double m, double v, double x){
int int_m;
double out;
if (m != floor(m))
{
return CUDART_NAN;
}
int_m = (int) m;
out = lpmv(v, int_m, x);
// should raise an overflow warning here on INF
return out;
}
"""
lpmv = _core.create_ufunc(
"cupyx_scipy_lpmv",
("fff->f", "ddd->d"),
"out0 = out0_type(pmv_wrap(in0, in1, in2));",
preamble=lpmv_definition,
doc="""Associated Legendre function of integer order and real degree.
.. seealso:: :meth:`scipy.special.lpmv`
""",
)
from cupy import _core
j0 = _core.create_ufunc('cupyx_scipy_special_j0', ('f->f', 'd->d'),
'out0 = j0(in0)',
doc='''Bessel function of the first kind of order 0.
.. seealso:: :meth:`scipy.special.j0`
''')
j1 = _core.create_ufunc('cupyx_scipy_special_j1', ('f->f', 'd->d'),
'out0 = j1(in0)',
doc='''Bessel function of the first kind of order 1.
.. seealso:: :meth:`scipy.special.j1`
''')
y0 = _core.create_ufunc('cupyx_scipy_special_y0', ('f->f', 'd->d'),
'out0 = y0(in0)',
doc='''Bessel function of the second kind of order 0.
.. seealso:: :meth:`scipy.special.y0`
''')
y1 = _core.create_ufunc('cupyx_scipy_special_y1', ('f->f', 'd->d'),
'out0 = y1(in0)',
doc='''Bessel function of the second kind of order 1.
.. seealso:: :meth:`scipy.special.y1`
from cupy import _core
gammaln = _core.create_ufunc(
'cupyx_scipy_special_gammaln', ('f->f', 'd->d'),
'''
if (isinf(in0) && in0 < 0) {
out0 = -1.0 / 0.0;
} else {
out0 = lgamma(in0);
}
''',
doc="""Logarithm of the absolute value of the Gamma function.
Args:
x (cupy.ndarray): Values on the real line at which to compute
``gammaln``.
Returns:
cupy.ndarray: Values of ``gammaln`` at x.
.. seealso:: :data:`scipy.special.gammaln`
""")
return a.ravel().nonzero()[0]
_where_ufunc = _core.create_ufunc(
'cupy_where',
(
'???->?',
'?bb->b',
'?BB->B',
'?hh->h',
'?HH->H',
'?ii->i',
'?II->I',
'?ll->l',
'?LL->L',
'?qq->q',
'?QQ->Q',
'?ee->e',
'?ff->f',
# On CUDA 6.5 these combinations don't work correctly (on CUDA >=7.0, it
# works).
# See issue #551.
'?hd->d',
'?Hd->d',
'?dd->d',
'?FF->F',
'?DD->D'),
'out0 = in0 ? in1 : in2')
def where(condition, x=None, y=None):
def _create_float_test_ufunc(name, doc):
return _core.create_ufunc(
'cupy_' + name,
('e->?', 'f->?', 'd->?', 'F->?', 'D->?',
), 'out0 = %s(in0)' % name,
doc=doc)
return 0;
} else {
double u = delta;
double v = r / delta;
return u * u * (sqrt(1 + v * v) - 1);
}
}
'''
entr = _core.create_ufunc(
'cupyx_scipy_special_entr', ('f->f', 'd->d'),
'out0 = out0_type(entr(in0));',
preamble=_float_preamble,
doc='''Elementwise function for computing entropy.
.. seealso:: :meth:`scipy.special.entr`
''')
kl_div = _core.create_ufunc(
'cupyx_scipy_special_kl_div', ('ff->f', 'dd->d'),
'out0 = out0_type(kl_div(in0, in1));',
preamble=_float_preamble,
doc='''Elementwise function for computing Kullback-Leibler divergence.
.. seealso:: :meth:`scipy.special.kl_div`
''')
out0 = tgamma(in0);
}
"""
# Also define a standalone Gamma device function for internal use in other code
# like beta, betaln, etc.
gamma_definition = f"""
__noinline__ __device__ double Gamma(double in0)
{{
double out0;
{_gamma_body}
return out0;
}}
"""
gamma = _core.create_ufunc(
'cupyx_scipy_gamma', ('f->f', 'd->d'),
_gamma_body,
doc="""Gamma function.
Args:
z (cupy.ndarray): The input of gamma function.
Returns:
cupy.ndarray: Computed value of gamma function.
.. seealso:: :data:`scipy.special.gamma`
""")
Cephes Math Library, Release 2.3: March, 1995
Copyright 1984, 1995 by Stephen L. Moshier
"""
from cupy import _core
from cupyx.scipy.special._beta import incbet_preamble, incbi_preamble
from cupyx.scipy.special._gammainc import _igam_preamble, _igami_preamble
# Normal distribution functions
ndtr = _core.create_ufunc(
'cupyx_scipy_special_ndtr',
(('f->f', 'out0 = normcdff(in0)'), 'd->d'),
'out0 = normcdf(in0)',
doc='''Cumulative distribution function of normal distribution.
.. seealso:: :data:`scipy.special.ndtr`
''')
ndtri = _core.create_ufunc(
'cupyx_scipy_special_ndtri',
(('f->f', 'out0 = normcdfinvf(in0)'), 'd->d'),
'out0 = normcdfinv(in0)',
doc='''Inverse of the cumulative distribution function of the standard
normal distribution.
.. seealso:: :data:`scipy.special.ndtri`
''')
.. seealso:: :data:`numpy.exp`
''')
expm1 = ufunc.create_math_ufunc(
'expm1', 1, 'cupy_expm1', '''Computes ``exp(x) - 1`` elementwise.
.. seealso:: :data:`numpy.expm1`
''')
exp2 = _core.create_ufunc('cupy_exp2',
('e->e', 'f->f', 'd->d', 'F->F', 'D->D'),
'out0 = pow(in0_type(2), in0)',
doc='''Elementwise exponentiation with base 2.
.. seealso:: :data:`numpy.exp2`
''')
log = ufunc.create_math_ufunc(
'log', 1, 'cupy_log', '''Elementwise natural logarithm function.
.. seealso:: :data:`numpy.log`
''')
log10 = ufunc.create_math_ufunc(
'log10', 1, 'cupy_log10', '''Elementwise common logarithm function.
.. seealso:: :data:`numpy.log10`
import numpy
import cupy
from cupy import _core
from cupy._logic import content
_is_close = _core.create_ufunc(
'cupy_is_close',
('eeee?->?', 'ffff?->?', 'dddd?->?'),
'''
bool equal_nan = in4;
if (isfinite(in0) && isfinite(in1)) {
out0 = fabs(in0 - in1) <= in3 + in2 * fabs(in1);
} else if (equal_nan) {
out0 = (in0 == in1) || (isnan(in0) && isnan(in1));
} else {
out0 = (in0 == in1);
}
'''
)
# Note that in cupy/_core/include/cupy/complex.cuh, we already got isfinite and
# isnan working for complex numbers, so just replace fabs above by abs (from
# thrust) and we are ready to go
_is_close_complex = _core.create_ufunc(
'cupy_is_close_complex',
('FFff?->?', 'DDdd?->?'),
'''
bool equal_nan = in4;
if (isfinite(in0) && isfinite(in1)) {
return (ans);
}
''')
expn = _core.create_ufunc(
'cupyx_scipy_special_expn',
('ff->f', 'dd->d'),
'out0 = expn(in0, in1)',
preamble=expn_definition,
doc="""Generalized exponential integral En.
Parameters
----------
n : cupy.ndarray
Non-negative integers
x : cupy.ndarray
Real argument
Returns
-------
y : scalar or cupy.ndarray
Values of the generalized exponential integral
See Also
--------
:func:`scipy.special.expn`
""",
)
t = a * b / A[i];
s = s + t;
t = fabs(t / s);
if (t < MACHEP){
return s;
}
k += 1.0;
a *= x + k;
b /= w;
k += 1.0;
}
return s;
}
'''
zeta = _core.create_ufunc('cupyx_scipy_special_zeta', ('ff->f', 'dd->d'),
'out0 = zeta(in0, in1)',
preamble=zeta_definition,
doc="""Hurwitz zeta function.
Args:
x (cupy.ndarray): Input data, must be real.
q (cupy.ndarray): Input data, must be real.
Returns:
cupy.ndarray: Values of zeta(x, q).
.. seealso:: :data:`scipy.special.zeta`
""")
from cupy import _core
erf = _core.create_ufunc('cupyx_scipy_erf', ('f->f', 'd->d'),
'out0 = erf(in0)',
doc='''Error function.
.. seealso:: :meth:`scipy.special.erf`
''')
erfc = _core.create_ufunc('cupyx_scipy_erfc', ('f->f', 'd->d'),
'out0 = erfc(in0)',
doc='''Complementary error function.
.. seealso:: :meth:`scipy.special.erfc`
''')
erfcx = _core.create_ufunc('cupyx_scipy_erfcx', ('f->f', 'd->d'),
'out0 = erfcx(in0)',
doc='''Scaled complementary error function.
.. seealso:: :meth:`scipy.special.erfcx`
''')
erfinv = _core.create_ufunc('cupyx_scipy_erfinv', ('f->f', 'd->d'),
'out0 = erfinv(in0);',
doc='''Inverse function of error function.
.. seealso:: :meth:`scipy.special.erfinv`
r = in0 % in1;
in0 = in1;
in1 = r;
}
if (in0 < 0)
return -in0;
return in0;
}
'''
gcd = _core.create_ufunc(
'cupy_gcd',
(('??->?', _negative_gcd_error),
'bb->b', 'BB->B', 'hh->h', 'HH->H', 'ii->i', 'II->I', 'll->l',
'LL->L', 'qq->q', 'QQ->Q'),
'out0 = gcd(in0, in1)',
preamble=_gcd_preamble,
doc='''Computes gcd of ``x1`` and ``x2`` elementwise.
.. seealso:: :data:`numpy.gcd`
''')
_lcm_preamble = _gcd_preamble + '''
template <typename T> inline __device__ T lcm(T in0, T in1) {
T r = gcd(in0, in1);
if (r == 0)
return 0;
r = in0 / r * in1;
if (r < 0)
return -r;
return r;
"""
beta = _core.create_ufunc(
"cupyx_scipy_beta",
("ff->f", "dd->d"),
"out0 = out0_type(beta(in0, in1));",
preamble=(beta_preamble + gamma_definition + polevl_definition +
p1evl_definition + lgam_sgn_definition + lbeta_symp_definition +
beta_definition),
doc="""Beta function.
Parameters
----------
a, b : cupy.ndarray
Real-valued arguments
out : cupy.ndarray, optional
Optional output array for the function result
Returns
-------
scalar or ndarray
Value of the beta function
See Also
--------
:func:`scipy.special.beta`
""",
)
betaln = _core.create_ufunc(
from cupy import _core
gamma = _core.create_ufunc('cupyx_scipy_gamma', ('f->f', 'd->d'),
'''
if (isinf(in0) && in0 < 0) {
out0 = -1.0 / 0.0;
} else if (in0 < 0. && in0 == floor(in0)) {
out0 = 1.0 / 0.0;
} else {
out0 = tgamma(in0);
}
''',
doc="""Gamma function.
Args:
z (cupy.ndarray): The input of gamma function.
Returns:
cupy.ndarray: Computed value of gamma function.
.. seealso:: :data:`scipy.special.gamma`
""")
+ m*(m-1)*(m-2)*(3*m-1)/(24*a*a)
+ m*m*(m-1)*(m-1)*(m-2)*(m-3)/(48*a*a*a)
);
}
/* Check for infinity */
if (is_nonpos_int(a + m) && !is_nonpos_int(a) && a + m != m) {
return CUDART_INF;
}
/* Check for zero */
if (!is_nonpos_int(a + m) && is_nonpos_int(a)) {
return 0;
}
return r * exp(lgam(a + m) - lgam(a)) * gammasgn(a + m) * gammasgn(a);
}
""")
poch = _core.create_ufunc(
"cupyx_scipy_poch",
("ff->f", "dd->d"),
"out0 = out0_type(poch(in0, in1));",
preamble=poch_definition,
doc="""Elementwise function for scipy.special.poch (Pochhammer symbol)
.. seealso:: :meth:`scipy.special.poch`
""",
)
from cupy import _core
from cupy._core import fusion
add = _core.add
reciprocal = _core.create_ufunc(
'cupy_reciprocal',
('b', 'B', 'h', 'H', 'i', 'I', 'l', 'L', 'q', 'Q', ('e', 'out0 = 1 / in0'),
('f', 'out0 = 1 / in0'), ('d', 'out0 = 1 / in0'),
('F', 'out0 = in0_type(1) / in0'), ('D', 'out0 = in0_type(1) / in0')),
'out0 = in0 == 0 ? 0 : (1 / in0)',
doc='''Computes ``1 / x`` elementwise.
.. seealso:: :data:`numpy.reciprocal`
''')
positive = _core.positive
negative = _core.negative
conjugate = _core.conjugate
# cupy.real is not a ufunc because it returns a view.
# The ufunc implementation is used by fusion.
_real_ufunc = _core.create_ufunc(
'cupy_real',
('?->?', 'b->b', 'B->B', 'h->h', 'H->H', 'i->i', 'I->I', 'l->l', 'L->L',
'q->q', 'Q->Q', 'e->e', 'f->f', 'd->d', ('F->f', 'out0 = in0.real()'),
('D->d', 'out0 = in0.real()')),
'out0 = in0',
from cupy import _core
j0 = _core.create_ufunc(
'cupyx_scipy_special_j0', ('f->f', 'd->d'),
'out0 = j0(in0)',
doc='''Bessel function of the first kind of order 0.
.. seealso:: :meth:`scipy.special.j0`
''')
j1 = _core.create_ufunc(
'cupyx_scipy_special_j1', ('f->f', 'd->d'),
'out0 = j1(in0)',
doc='''Bessel function of the first kind of order 1.
.. seealso:: :meth:`scipy.special.j1`
''')
y0 = _core.create_ufunc(
'cupyx_scipy_special_y0', ('f->f', 'd->d'),
'out0 = y0(in0)',
doc='''Bessel function of the second kind of order 0.
.. seealso:: :meth:`scipy.special.y0`
''')
import ast
import numpy
from cupy._logic import ops
from cupy._math import arithmetic
from cupy._logic import comparison
from cupy._binary import elementwise
from cupy import _core
from cupyx.jit import _types
_numpy_scalar_true_divide = _core.create_ufunc(
'numpy_scalar_true_divide',
('??->d', '?i->d', 'i?->d', 'bb->f', 'bi->d', 'BB->f', 'Bi->d', 'hh->f',
'HH->f', 'ii->d', 'II->d', 'll->d', 'LL->d', 'qq->d', 'QQ->d', 'ee->e',
'ff->f', 'dd->d', 'FF->F', 'DD->D'),
'out0 = (out0_type)in0 / (out0_type)in1',
)
_numpy_scalar_invert = _core.create_ufunc(
'numpy_scalar_invert',
('?->?', 'b->b', 'B->B', 'h->h', 'H->H', 'i->i', 'I->I', 'l->l', 'L->L',
'q->q', 'Q->Q'),
'out0 = ~in0',
)
_numpy_scalar_logical_not = _core.create_ufunc(
'numpy_scalar_logical_not',
('?->?', 'b->?', 'B->?', 'h->?', 'H->?', 'i->?', 'I->?', 'l->?', 'L->?',
'q->?', 'Q->?', 'e->?', 'f->?', 'd->?',
.. seealso:: :data:`numpy.hypot`
''')
arctan2 = ufunc.create_math_ufunc(
'atan2', 2, 'cupy_arctan2',
'''Elementwise inverse-tangent of the ratio of two arrays.
.. seealso:: :data:`numpy.arctan2`
''')
deg2rad = _core.create_ufunc(
'cupy_deg2rad', ('e->e', 'f->f', 'd->d'),
'out0 = in0 * (out0_type)(M_PI / 180)',
doc='''Converts angles from degrees to radians elementwise.
.. seealso:: :data:`numpy.deg2rad`, :data:`numpy.radians`
''')
rad2deg = _core.create_ufunc(
'cupy_rad2deg', ('e->e', 'f->f', 'd->d'),
'out0 = in0 * (out0_type)(180 / M_PI)',
doc='''Converts angles from radians to degrees elementwise.
.. seealso:: :data:`numpy.rad2deg`, :data:`numpy.degrees`
''')
def unwrap(p, discont=None, axis=-1, *, period=2 * numpy.pi):
import numpy
import numpy.typing as npt
import operator
import cupy
from cupy._logic import ops
from cupy._math import arithmetic
from cupy._logic import comparison
from cupy._binary import elementwise
from cupy import _core
from cupyx.jit import _cuda_types
_numpy_scalar_invert = _core.create_ufunc(
'numpy_scalar_invert',
('?->?', 'b->b', 'B->B', 'h->h', 'H->H', 'i->i', 'I->I', 'l->l', 'L->L',
'q->q', 'Q->Q'),
'out0 = ~in0',
)
_numpy_scalar_logical_not = _core.create_ufunc(
'numpy_scalar_logical_not',
('?->?', 'b->?', 'B->?', 'h->?', 'H->?', 'i->?', 'I->?', 'l->?', 'L->?',
'q->?', 'Q->?', 'e->?', 'f->?', 'd->?',
('F->?', 'out0 = !in0.real() && !in0.imag()'),
('D->?', 'out0 = !in0.real() && !in0.imag()')),
'out0 = !in0',
)
_scalar_lt = _core.create_comparison('scalar_less', '<')
_scalar_lte = _core.create_comparison('scalar_less', '<=')
_scalar_gt = _core.create_comparison('scalar_less', '>')
from cupy import _core
ndtr = _core.create_ufunc(
'cupyx_scipy_ndtr', ('f->f', 'd->d'),
'out0 = normcdf(in0)',
doc='''Cumulative distribution function of normal distribution.
.. seealso:: :meth:`scipy.special.ndtr`
''')
static __device__ T logit(T x) {
x /= 1 - x;
return log(x);
}
"""
logit = _core.create_ufunc(
'cupy_logit',
('e->f', 'f->f', 'd->d'),
'out0 = logit(in0)',
preamble=logit_definition,
doc='''Logit function.
Args:
x (cupy.ndarray): input data
Returns:
cupy.ndarray: values of logit(x)
.. seealso:: :data:`scipy.special.logit`
''')
expit_definition = """
template <typename T>
static __device__ T expit(T x) {
return 1 / (1 + exp(-x));
}
from cupy import _core
from cupy._math import ufunc
from cupy.cuda import runtime
signbit = _core.create_ufunc(
'cupy_signbit', ('e->?', 'f->?', 'd->?'),
'out0 = signbit(in0)',
doc='''Tests elementwise if the sign bit is set (i.e. less than zero).
.. seealso:: :data:`numpy.signbit`
''')
copysign = ufunc.create_math_ufunc(
'copysign', 2, 'cupy_copysign',
'''Returns the first argument with the sign bit of the second elementwise.
.. seealso:: :data:`numpy.copysign`
''')
ldexp = _core.create_ufunc(
'cupy_ldexp', ('ei->e', 'fi->f', 'el->e', 'fl->f', 'di->d', 'dq->d'),
'out0 = ldexp(in0, in1)',
doc='''Computes ``x1 * 2 ** x2`` elementwise.
.. seealso:: :data:`numpy.ldexp`
''')
# HIP supports frexpf but not frexp ...