import functools
import operator
from collections.abc import Callable
from dataclasses import dataclass
from types import ModuleType
import numpy as np
from scipy._lib._array_api import (
array_namespace, scipy_namespace_for, is_numpy, is_dask, is_marray,
xp_promote, xp_capabilities, SCIPY_ARRAY_API, get_native_namespace_name
)
import scipy._lib.array_api_extra as xpx
from . import _basic
from . import _spfun_stats
from . import _ufuncs
@dataclass
class _FuncInfo:
# NumPy-only function. IT MUST BE ELEMENTWISE.
func: Callable
# Number of arguments, not counting out=
# This is for testing purposes only, due to the fact that
# inspect.signature() just returns *args for ufuncs.
n_args: int
# @xp_capabilities decorator, for the purpose of
# documentation and unit testing. Omit to indicate
# full support for all backends.
xp_capabilities: Callable[[Callable], Callable] | None = None
# Generic implementation to fall back on if there is no native dispatch
# available. This is a function that accepts (main namespace, scipy namespace)
# and returns the final callable, or None if not available.
generic_impl: Callable[
[ModuleType, ModuleType | None], Callable | None
] | None = None
# Handle case where a backend uses an alternative name for a function.
# Should map backend names to alternative function names.
alt_names_map: dict[str, str] | None = None
# Some functions only take integer arrays for some arguments.
int_only: tuple[bool] | None = None
# For testing purposes, whether tests should only use positive values
# for some arguments. If bool and equal to True, restrict to positive
# values for all arguments. To restrict only some arguments to positive
# values, pass a tuple of bool of the same length as the number of
# arguments, the ith entry in the tuple controls positive_only for
# the ith argument. To make backend specific choices for positive_only,
# pass in a dict mapping backend names to bool or tuple[bool].
positive_only: bool | tuple[bool] | dict[str, tuple[bool]] = False
# Some special functions are not ufuncs and ufunc-specific tests
# should not be applied to these.
is_ufunc: bool = True
# Some non-ufunc special functions take only Python ints for some arguments.
# If so, python_int_only should be a tuple of the same length as the number
# of arguments,with value True if the corresponding argument needs to be a
# Python int.
# Can also take a dict mapping backends to such tuples if an argument being
# Python int only is backend specific.
python_int_only: dict[str, tuple[bool]] | tuple[bool] | None = None
# Some functions which seem to be scalar also accept 0d arrays.
scalar_or_0d_only: dict[str, tuple[bool]] | tuple[bool] | None = None
# Some functions may not work well with very large integer valued arguments.
test_large_ints: bool = True
# Some non-ufunc special functions don't decay 0d arrays to scalar.
produces_0d: bool = False
# Whether or not uses native PyTorch or falls back to NumPy/SciPy. This
# is needed because in PyTorch, the default dtype affects promotion
# rules when mixing integer and floating dtypes, so relying on a
# NumPy/SciPy fallback when the default dtype is other than float64 can lead
# to float64 output when native PyTorch would have e.g. float32 output. This
# must be accounted for in tests. Not putting this in xp_capabilities for now,
# but in the future I think it's likely we may want to add a warning to
# xp_capabilities when not using native PyTorch on CPU.
torch_native: bool = True
@property
def name(self):
return self.func.__name__
# These are needed by @lru_cache below
def __hash__(self):
return hash(self.func)
def __eq__(self, other):
return isinstance(other, _FuncInfo) and self.func == other.func
@property
def wrapper(self):
if self.name in globals():
# Already initialised. We are likely in a unit test.
# Return function potentially overridden by xpx.testing.lazy_xp_function.
import scipy.special
return getattr(scipy.special, self.name)
if SCIPY_ARRAY_API:
@functools.wraps(self.func)
def wrapped(*args, **kwargs):
xp = array_namespace(*args)
return self._wrapper_for(xp)(*args, **kwargs)
# Allow pickling the function. Normally this is done by @wraps,
# but in this case it doesn't work because self.func is a ufunc.
wrapped.__module__ = "scipy.special"
wrapped.__qualname__ = self.name
func = wrapped
else:
func = self.func
capabilities = self.xp_capabilities or xp_capabilities()
# In order to retain a naked ufunc when SCIPY_ARRAY_API is
# disabled, xp_capabilities must apply its changes in place.
cap_func = capabilities(func)
assert cap_func is func
return func
@functools.lru_cache(1000)
def _wrapper_for(self, xp):
if is_numpy(xp):
return self.func
# If a native implementation is available, use that
spx = scipy_namespace_for(xp)
f = _get_native_func(xp, spx, self.name, alt_names_map=self.alt_names_map)
if f is not None:
return f
# If generic Array API implementation is available, use that
if self.generic_impl is not None:
f = self.generic_impl(xp, spx)
if f is not None:
return f
if is_marray(xp):
# Unwrap the array, apply the function on the wrapped namespace,
# and then re-wrap it.
# IMPORTANT: this only works because all functions in this module
# are elementwise. Otherwise, we would not be able to define a
# general rule for mask propagation.
_f = globals()[self.name] # Allow nested wrapping
def f(*args, _f=_f, xp=xp, **kwargs):
data_args = [getattr(arg, 'data', arg) for arg in args]
out = _f(*data_args, **kwargs)
mask = functools.reduce(operator.or_,
(getattr(arg, 'mask', False) for arg in args))
return xp.asarray(out, mask=mask)
return f
if is_dask(xp):
# Apply the function to each block of the Dask array.
# IMPORTANT: map_blocks works only because all functions in this module
# are elementwise. It would be a grave mistake to apply this to gufuncs
# or any other function with reductions, as they would change their
# output depending on chunking!
_f = globals()[self.name] # Allow nested wrapping
def f(*args, _f=_f, xp=xp, **kwargs):
# Hide dtype kwarg from map_blocks
return xp.map_blocks(functools.partial(_f, **kwargs), *args)
return f
# As a final resort, use the NumPy/SciPy implementation
_f = self.func
def f(*args, _f=_f, xp=xp, **kwargs):
# TODO use xpx.lazy_apply to add jax.jit support
# (but dtype propagation can be non-trivial)
args = [np.asarray(arg) for arg in args]
out = _f(*args, **kwargs)
return xp.asarray(out)
return f
def _get_native_func(xp, spx, f_name, *, alt_names_map=None):
if alt_names_map is None:
alt_names_map = {}
f_name = alt_names_map.get(get_native_namespace_name(xp), f_name)
f = getattr(spx.special, f_name, None) if spx else None
if f is None and hasattr(xp, 'special'):
# Currently dead branch, in anticipation of 'special' Array API extension
# https://github.com/data-apis/array-api/issues/725
f = getattr(xp.special, f_name, None)
return f
def _rel_entr(xp, spx):
def __rel_entr(x, y, *, xp=xp):
# https://github.com/data-apis/array-api-extra/issues/160
mxp = array_namespace(x._meta, y._meta) if is_dask(xp) else xp
x, y = xp_promote(x, y, broadcast=True, force_floating=True, xp=xp)
xy_pos = (x > 0) & (y > 0)
xy_inf = xp.isinf(x) & xp.isinf(y)
res = xpx.apply_where(
xy_pos & ~xy_inf,
(x, y),
# Note: for very large x, this can overflow.
lambda x, y: x * (mxp.log(x) - mxp.log(y)),
fill_value=xp.inf
)
res = xpx.at(res)[(x == 0) & (y >= 0)].set(0)
res = xpx.at(res)[xp.isnan(x) | xp.isnan(y) | (xy_pos & xy_inf)].set(xp.nan)
return res
return __rel_entr
def _xlogy(xp, spx):
def __xlogy(x, y, *, xp=xp):
x, y = xp_promote(x, y, force_floating=True, xp=xp)
with np.errstate(divide='ignore', invalid='ignore'):
temp = x * xp.log(y)
return xp.where(x == 0., 0., temp)
return __xlogy
def _chdtr(xp, spx):
# The difference between this and just using `gammainc`
# defined by `get_array_special_func` is that if `gammainc`
# isn't found, we don't want to use the SciPy version; we'll
# return None here and use the SciPy version of `chdtr`.
gammainc = _get_native_func(xp, spx, 'gammainc')
if gammainc is None:
return None
def __chdtr(v, x):
res = gammainc(v / 2, x / 2) # this is almost all we need
# The rest can be removed when google/jax#20507 is resolved
mask = (v == 0) & (x > 0) # JAX returns NaN
res = xp.where(mask, 1., res)
mask = xp.isinf(v) & xp.isinf(x) # JAX returns 1.0
return xp.where(mask, xp.nan, res)
return __chdtr
def _chdtrc(xp, spx):
# The difference between this and just using `gammaincc`
# defined by `get_array_special_func` is that if `gammaincc`
# isn't found, we don't want to use the SciPy version; we'll
# return None here and use the SciPy version of `chdtrc`.
gammaincc = _get_native_func(xp, spx, 'gammaincc')
if gammaincc is None:
return None
def __chdtrc(v, x):
res = xp.where(x >= 0, gammaincc(v/2, x/2), 1)
i_nan = ((x == 0) & (v == 0)) | xp.isnan(x) | xp.isnan(v) | (v <= 0)
res = xp.where(i_nan, xp.nan, res)
return res
return __chdtrc
def _betaincc(xp, spx):
betainc = _get_native_func(xp, spx, 'betainc')
if betainc is None:
return None
def __betaincc(a, b, x):
# not perfect; might want to just rely on SciPy
return betainc(b, a, 1-x)
return __betaincc
def _stdtr(xp, spx):
betainc = _get_native_func(xp, spx, 'betainc')
if betainc is None:
return None
def __stdtr(df, t):
x = df / (t ** 2 + df)
tail = betainc(df / 2, 0.5, x) / 2
return xp.where(t < 0, tail, 1 - tail)
return __stdtr
def _stdtrit(xp, spx):
# Need either native stdtr or native betainc
stdtr = _get_native_func(xp, spx, 'stdtr') or _stdtr(xp, spx)
# If betainc is not defined, the root-finding would be done with `xp`
# despite `stdtr` being evaluated with SciPy/NumPy `stdtr`. Save the
# conversions: in this case, just evaluate `stdtrit` with SciPy/NumPy.
if stdtr is None:
return None
from scipy.optimize.elementwise import bracket_root, find_root
def __stdtrit(df, p):
def fun(t, df, p): return stdtr(df, t) - p
res_bracket = bracket_root(fun, xp.zeros_like(p), args=(df, p))
res_root = find_root(fun, res_bracket.bracket, args=(df, p))
return res_root.x
return __stdtrit
# Inventory of automatically dispatched functions
# IMPORTANT: these must all be **elementwise** functions!
# PyTorch doesn't implement `betainc`.
# On torch CPU we can fall back to NumPy, but on GPU it won't work.
_needs_betainc = xp_capabilities(cpu_only=True, exceptions=["jax.numpy", "cupy"])
_special_funcs = (
_FuncInfo(
_ufuncs.bdtr, 3,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
int_only=(False, True, False), torch_native=False,
),
_FuncInfo(
_ufuncs.bdtrc, 3,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
int_only=(False, True, False), torch_native=False,
),
_FuncInfo(
_ufuncs.bdtri, 3,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
int_only=(False, True, False), torch_native=False,
),
_FuncInfo(_ufuncs.betainc, 3, _needs_betainc, torch_native=False),
_FuncInfo(_ufuncs.betaincc, 3, _needs_betainc, generic_impl=_betaincc,
torch_native=False),
_FuncInfo(
_ufuncs.betaincinv, 3,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
test_large_ints=False, positive_only=True, torch_native=False,
),
_FuncInfo(
_ufuncs.betaln, 2,
xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"]),
# For betaln, nan mismatches can occur at negative integer a or b of
# sufficiently large magnitude.
positive_only={"jax.numpy": True}, torch_native=False,
),
_FuncInfo(
_ufuncs.binom, 2,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(
_ufuncs.boxcox, 2,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(
_ufuncs.boxcox1p, 2,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(
_ufuncs.cbrt, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(_ufuncs.chdtr, 2, generic_impl=_chdtr),
_FuncInfo(_ufuncs.chdtrc, 2, generic_impl=_chdtrc,
# scipy/scipy#20972
positive_only={"cupy": True, "jax.numpy": True, "torch": True}),
_FuncInfo(
_ufuncs.chdtri, 2,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(
_ufuncs.cosdg, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
test_large_ints=False, torch_native=False,
),
_FuncInfo(
_ufuncs.cosm1, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(
_ufuncs.cotdg, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(
_ufuncs.ellipk, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(
_ufuncs.ellipkm1, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(_ufuncs.entr, 1),
_FuncInfo(_ufuncs.erf, 1),
_FuncInfo(_ufuncs.erfc, 1),
_FuncInfo(
_ufuncs.erfcx, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy", "torch"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(_ufuncs.erfinv, 1),
_FuncInfo(
_ufuncs.exp1, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(
_ufuncs.exp10, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(
_ufuncs.exp2, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(
_ufuncs.exprel, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(
_ufuncs.expi, 1,
xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"]),
torch_native=False,
),
_FuncInfo(_ufuncs.expit, 1),
_FuncInfo(
_ufuncs.expn, 2,
xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"]),
# Inconsistent behavior for negative n. expn is not defined here without
# taking analytic continuation.
positive_only=True,
int_only=(True, False), test_large_ints=False,
torch_native=False,
),
_FuncInfo(
_ufuncs.fdtr, 3,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(
_ufuncs.fdtrc, 3,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(
_ufuncs.fdtri, 3,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(
_ufuncs.gamma, 1,
xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"]),
torch_native=False,
),
_FuncInfo(_ufuncs.gammainc, 2),
_FuncInfo(
_ufuncs.gammaincc, 2,
# google/jax#20699
positive_only={"jax.numpy": True},
),
_FuncInfo(
_ufuncs.gammainccinv, 2,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(
_ufuncs.gammaincinv, 2,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(_ufuncs.gammaln, 1),
_FuncInfo(
_ufuncs.gammasgn, 1,
xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"]),
torch_native=False,
),
_FuncInfo(
_ufuncs.gdtr, 3,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(
_ufuncs.gdtrc, 3,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(
_ufuncs.huber, 2,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(
_ufuncs.hyp1f1, 3,
xp_capabilities(cpu_only=True, exceptions=["jax.numpy"]),
positive_only={"jax.numpy": True}, test_large_ints=False,
torch_native=False,
),
# Comment out when jax>=0.6.1 is available in Conda for CI.
# (or add version requirements to xp_capabilities).
# _FuncInfo(
# _ufuncs.hyp2f1, 4,
# xp_capabilities(cpu_only=True, exceptions=["jax.numpy"]),
# positive_only={"jax.numpy": True}, test_large_ints=False,
# torch_native=False,
# ),
_FuncInfo(
_ufuncs.inv_boxcox, 2,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(
_ufuncs.inv_boxcox1p, 2,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(_ufuncs.i0, 1),
_FuncInfo(_ufuncs.i0e, 1),
_FuncInfo(_ufuncs.i1, 1),
_FuncInfo(_ufuncs.i1e, 1),
_FuncInfo(
_ufuncs.j0, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy", "torch"],
jax_jit=False,
),
alt_names_map={"torch": "bessel_j0"}, test_large_ints=False,
),
_FuncInfo(
_ufuncs.j1, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy", "torch"],
jax_jit=False,
),
alt_names_map={"torch": "bessel_j1"}, test_large_ints=False,
),
_FuncInfo(
_ufuncs.k0, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy", "torch"],
jax_jit=False,
),
alt_names_map={"torch": "modified_bessel_k0"},
),
_FuncInfo(
_ufuncs.k0e, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy", "torch"],
jax_jit=False,
),
alt_names_map={"torch": "scaled_modified_bessel_k0"},
test_large_ints=False,
),
_FuncInfo(
_ufuncs.k1, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy", "torch"],
jax_jit=False,
),
alt_names_map={"torch": "modified_bessel_k1"},
),
_FuncInfo(
_ufuncs.k1e, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy", "torch"],
jax_jit=False,
),
alt_names_map={"torch": "scaled_modified_bessel_k1"},
test_large_ints=False),
_FuncInfo(
_ufuncs.kl_div, 2,
xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"]),
torch_native=False,
),
_FuncInfo(_ufuncs.log_ndtr, 1),
_FuncInfo(
_ufuncs.loggamma, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(_ufuncs.logit, 1),
_FuncInfo(
_ufuncs.lpmv, 3,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
test_large_ints=False,
),
_FuncInfo(
_spfun_stats.multigammaln, 2,
is_ufunc=False,
python_int_only={
"cupy": [False, True],
"jax.numpy": [False, True],
"torch": [False, True],
},
scalar_or_0d_only={
"array_api_strict": [False, True],
"numpy": [False, True],
"dask.array": [False, True],
"marray": [False, True],
},
int_only=(False, True), test_large_ints=False,
positive_only=True, torch_native=False,
),
_FuncInfo(
_ufuncs.nbdtr, 3,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
int_only=(True, True, False), positive_only=True,
torch_native=False,
),
_FuncInfo(
_ufuncs.nbdtrc, 3,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
int_only=(True, True, False), positive_only=True,
torch_native=False,
),
_FuncInfo(
_ufuncs.nbdtri, 3,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
int_only=(True, True, False), positive_only=True,
torch_native=False,
),
_FuncInfo(_ufuncs.ndtr, 1),
_FuncInfo(_ufuncs.ndtri, 1),
_FuncInfo(
_ufuncs.pdtr, 2,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
positive_only=True, torch_native=False,
),
_FuncInfo(
_ufuncs.pdtrc, 2,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
positive_only=True, torch_native=False,
),
_FuncInfo(
_ufuncs.pdtri, 2,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
int_only=(True, False), positive_only=True,
torch_native=False,
),
_FuncInfo(
_ufuncs.poch, 2,
xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"]),
test_large_ints=False, torch_native=False,
),
_FuncInfo(
_ufuncs.pseudo_huber, 2,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(
_basic.polygamma, 2, int_only=(True, False), is_ufunc=False,
scalar_or_0d_only={"torch": (True, False)}, produces_0d=True,
positive_only={"torch": (True, False), "jax.numpy": True},
test_large_ints=False,
),
_FuncInfo(_ufuncs.psi, 1, alt_names_map={"jax.numpy": "digamma"}),
_FuncInfo(
_ufuncs.radian, 3,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(_ufuncs.rel_entr, 2, generic_impl=_rel_entr),
_FuncInfo(
_ufuncs.rgamma, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
_FuncInfo(
_basic.sinc, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy", "torch"],
jax_jit=False,
),
is_ufunc=False,
),
_FuncInfo(
_ufuncs.sindg, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
test_large_ints=False, torch_native=False,
),
_FuncInfo(
_ufuncs.spence, 1,
xp_capabilities(cpu_only=True, exceptions=["jax.numpy"]),
torch_native=False,
),
_FuncInfo(_ufuncs.stdtr, 2, _needs_betainc, generic_impl=_stdtr,
torch_native=False),
_FuncInfo(
_ufuncs.stdtrit, 2,
xp_capabilities(
cpu_only=True, exceptions=["cupy"], # needs betainc
skip_backends=[("jax.numpy", "no scipy.optimize support")],
),
generic_impl=_stdtrit, torch_native=False,
),
_FuncInfo(
_ufuncs.tandg, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
test_large_ints=False, torch_native=False,
),
_FuncInfo(_ufuncs.xlog1py, 2),
_FuncInfo(_ufuncs.xlogy, 2, generic_impl=_xlogy),
_FuncInfo(
_ufuncs.y0, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy", "torch"],
jax_jit=False,
),
alt_names_map={"torch": "bessel_y0"}, test_large_ints=False,
),
_FuncInfo(
_ufuncs.y1, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy", "torch"],
jax_jit=False,
),
alt_names_map={"torch": "bessel_y1"}, test_large_ints=False,
),
_FuncInfo(
_ufuncs.yn, 2,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
positive_only={"cupy": (True, False)}, int_only=(True, False),
test_large_ints=False, torch_native=False,
),
_FuncInfo(
_basic.zeta, 2, is_ufunc=False,
positive_only={"jax.numpy": True, "torch": (True, False)},
test_large_ints=False,
),
_FuncInfo(
_ufuncs.zetac, 1,
xp_capabilities(
cpu_only=True, exceptions=["cupy"],
jax_jit=False,
),
torch_native=False,
),
)
# Override ufuncs.
# When SCIPY_ARRAY_API is disabled, this exclusively updates the docstrings in place
# and populates the xp_capabilities table, while retaining the original ufuncs.
globals().update({nfo.func.__name__: nfo.wrapper for nfo in _special_funcs})
# digamma is an alias for psi. Define here so it also has alternative backend
# support. Add noqa because the linter gets confused by the sneaky way psi
# is inserted into globals above.
digamma = psi # noqa: F821
__all__ = [nfo.func.__name__ for nfo in _special_funcs] + ["digamma"]