from functools import partial import pickle import pytest from hypothesis import given, strategies import hypothesis.extra.numpy as npst from packaging import version from scipy import special from scipy.special._support_alternative_backends import _special_funcs from scipy._lib._array_api_no_0d import xp_assert_close from scipy._lib._array_api import (is_cupy, is_dask, is_jax, is_torch, make_xp_pytest_param, make_xp_test_case, get_native_namespace_name) from scipy._lib.array_api_compat import numpy as np import scipy._lib.array_api_extra as xpx # Run all tests in this module in the Array API CI, including those without # the xp fixture pytestmark = pytest.mark.array_api_backends lazy_xp_modules = [special] def _skip_or_tweak_alternative_backends(xp, nfo, dtypes, int_only): """Skip tests for specific intersections of scipy.special functions vs. backends vs. dtypes vs. devices. Also suggest bespoke tweaks. Returns ------- positive_only : list[bool] Whether you should exclusively test positive inputs. dtypes : list[str] dtype strings 'float64', 'int32', 'int64', etc. with integer types mapped to the type of the NumPy default int. """ f_name = nfo.name if isinstance(nfo.positive_only, dict): positive_only = nfo.positive_only.get(get_native_namespace_name(xp), False) else: positive_only = nfo.positive_only if isinstance(positive_only, bool): positive_only = [positive_only]*nfo.n_args dtypes = [np.intp.__name__ if dtype == "intp" else dtype for dtype in dtypes] if f_name in {'betaincinv'} and is_cupy(xp): pytest.xfail("CuPy uses different convention for out of domain input.") if ( f_name == "sinc" and "float32" in dtypes and version.parse(np.__version__) < version.parse("2") ): pytest.xfail("https://github.com/numpy/numpy/issues/11204") if not any('int' in dtype for dtype in dtypes): return positive_only, dtypes # Integer-specific issues from this point onwards if f_name in {'gamma', 'gammasgn'} and is_cupy(xp): # CuPy has not yet updated gamma pole behavior to match # https://github.com/scipy/scipy/pull/21827. positive_only = [True] if f_name in {'poch'} and is_jax(xp): # Jax uses a different convention at gamma poles. positive_only = [True, True] if f_name == 'multigammaln': pytest.skip("multigammaln raises for out of domain inputs.") if ((is_torch(xp) and f_name in {'gammainc', 'gammaincc'}) or (is_cupy(xp) and f_name in {'stdtr', 'i0e', 'i1e'}) or (is_jax(xp) and f_name in {'stdtr', 'ndtr', 'ndtri', 'log_ndtr', 'hyp1f1', 'hyp2f1', 'spence', 'kl_div'}) ): pytest.skip(f"`{f_name}` does not support integer types") # int/float mismatched args support is sketchy if (any('float' in dtype for dtype in dtypes) and ((is_torch(xp) and f_name in ('rel_entr', 'xlogy', 'polygamma', 'zeta', 'xlog1py')) or (is_jax(xp) and f_name in ('gammainc', 'gammaincc', 'expn', 'rel_entr', 'xlogy', 'betaln', 'polygamma', 'zeta', 'poch', 'xlog1py'))) ): pytest.xfail("dtypes do not match") if (is_torch(xp) and xpx.default_dtype(xp) == xp.float32): # On PyTorch with float32 default dtype, all ints are promoted # to float32, but when falling back to NumPy/SciPy int64 is promoted # instead to float64. Integer only parameters essentially do not # participate in determination of the result type in PyTorch with # float32 default dtype, but will impact the output dtype as if # they were float64 when falling back to NumPy/SciPy. if not nfo.torch_native: pytest.xfail("dtypes do not match") return positive_only, dtypes @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") @pytest.mark.filterwarnings("ignore:overflow encountered:RuntimeWarning") @pytest.mark.parametrize('shapes', [[(0,)]*4, [tuple()]*4, [(10,)]*4, [(10,), (11, 1), (12, 1, 1), (13, 1, 1, 1)]]) @pytest.mark.parametrize('base_dtype', ['float32', 'float64', 'intp']) @pytest.mark.parametrize( 'func,nfo', [make_xp_pytest_param(i.wrapper, i) for i in _special_funcs]) def test_support_alternative_backends(xp, func, nfo, base_dtype, shapes): int_only = nfo.int_only if int_only is None: int_only = (False, ) * nfo.n_args dtypes = (base_dtype, ) * nfo.n_args else: dtypes = tuple( 'intp' if needs_int else base_dtype for needs_int in int_only ) positive_only, dtypes = _skip_or_tweak_alternative_backends( xp, nfo, dtypes, int_only ) dtypes_np = [getattr(np, dtype) for dtype in dtypes] dtypes_xp = [getattr(xp, dtype) for dtype in dtypes] shapes = shapes[:nfo.n_args] rng = np.random.default_rng(984254252920492019) args_np = [] # Handle cases where there's an argument which only takes scalar values. python_int_only = nfo.python_int_only if isinstance(python_int_only, dict): python_int_only = python_int_only.get(get_native_namespace_name(xp)) scalar_or_0d_only = nfo.scalar_or_0d_only if isinstance(scalar_or_0d_only, dict): scalar_or_0d_only = scalar_or_0d_only.get(get_native_namespace_name(xp)) test_large_ints = nfo.test_large_ints if isinstance(nfo.test_large_ints, dict): test_large_ints = test_large_ints.get(get_native_namespace_name(xp), False) if python_int_only is None: python_int_only = [False] * nfo.n_args if scalar_or_0d_only is None: scalar_or_0d_only = [False] * nfo.n_args no_shape = [ cond1 or cond2 for cond1, cond2 in zip(python_int_only, scalar_or_0d_only) ] shapes = [shape if not cond else None for shape, cond in zip(shapes, no_shape)] for dtype, dtype_np, shape, needs_python_int in zip( dtypes, dtypes_np, shapes, python_int_only ): if 'int' in dtype and test_large_ints: iinfo = np.iinfo(dtype_np) rand = partial(rng.integers, iinfo.min, iinfo.max + 1) elif 'int' in dtype: rand = partial(rng.integers, -20, 21) else: rand = rng.standard_normal val = rand(size=shape, dtype=dtype_np) if needs_python_int: # The logic above for determining shapes guarantees that # shape will be None in the above line when a Python int is required, # so this can safely be converted to an int. val = int(val) args_np.append(val) args_np = [ np.abs(arg) if cond else arg for arg, cond in zip(args_np, positive_only) ] args_xp = [ xp.asarray(arg, dtype=dtype_xp) if not needs_python_int else arg for arg, dtype_xp, needs_python_int in zip(args_np, dtypes_xp, python_int_only) ] args_np = [ np.asarray(arg, dtype=dtype_np) if not needs_python_int else arg for arg, dtype_np, needs_python_int in zip(args_np, dtypes_np, python_int_only) ] if is_dask(xp): # We're using map_blocks to dispatch the function to Dask. # This is the correct thing to do IF all tested functions are elementwise; # otherwise the output would change depending on chunking. # Try to trigger bugs related to having multiple chunks. args_xp = [arg.rechunk(5) for arg in args_xp] res = nfo.wrapper(*args_xp) # Also wrapped by lazy_xp_function ref = nfo.func(*args_np) # Unwrapped ufunc if ( is_torch(xp) and xpx.default_dtype(xp) == xp.float32 and "float64" not in dtypes ): # int64 promotes like float32 on torch with default dtype = float32 # cast reference if needed ref = np.float32(ref) # When dtype_np is integer, the output dtype can be float atol = 0 if ref.dtype.kind in 'iu' else 10 * np.finfo(ref.dtype).eps rtol = None if is_torch(xp) and func.__name__ == 'j1': # If we end up needing more function/backend specific tolerance # adjustments, this should be factored out properly. atol = 1e-7 rtol = 1e-5 xp_assert_close( res, xp.asarray(ref), rtol=rtol, atol=atol, check_0d=nfo.produces_0d ) @pytest.mark.parametrize( 'func, nfo', [make_xp_pytest_param(i.wrapper, i) for i in _special_funcs if i.n_args >= 2]) @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") def test_support_alternative_backends_mismatched_dtypes(xp, func, nfo): """Test mix-n-match of int and float arguments""" if func.__name__ in {'expn', 'polygamma', 'multigammaln', 'bdtr', 'bdtrc', 'bdtri', 'nbdtr', 'nbdtrc', 'nbdtri', 'pdtri'}: pytest.skip(f"dtypes for {func.__name__} make it a bad fit for this test.") dtypes = ['intp', 'float32', 'float64', 'float64'][:nfo.n_args] positive_only, dtypes = _skip_or_tweak_alternative_backends( xp, nfo, dtypes, (False,)*nfo.n_args ) dtypes_np = [getattr(np, dtype) for dtype in dtypes] dtypes_xp = [getattr(xp, dtype) for dtype in dtypes] rng = np.random.default_rng(984254252920492019) iinfo = np.iinfo(np.intp) if nfo.test_large_ints: randint = partial(rng.integers, iinfo.min, iinfo.max + 1) else: randint = partial(rng.integers, -20, 21) args_np = [ randint(size=1, dtype=np.int64), rng.standard_normal(size=1, dtype=np.float32), rng.standard_normal(size=1, dtype=np.float64), rng.standard_normal(size=1, dtype=np.float64), ][:nfo.n_args] args_np = [ np.abs(arg) if cond else arg for arg, cond in zip(args_np, positive_only) ] args_xp = [xp.asarray(arg, dtype=dtype_xp) for arg, dtype_xp in zip(args_np, dtypes_xp)] args_np = [np.asarray(arg, dtype=dtype_np) for arg, dtype_np in zip(args_np, dtypes_np)] res = nfo.wrapper(*args_xp) # Also wrapped by lazy_xp_function ref = nfo.func(*args_np) # Unwrapped ufunc if ( is_torch(xp) and xpx.default_dtype(xp) == xp.float32 and "float64" not in dtypes ): # int64 promotes like float32 on torch with default dtype = float32 # cast reference if needed ref = np.float32(ref) atol = 10 * np.finfo(ref.dtype).eps xp_assert_close(res, xp.asarray(ref), atol=atol) @pytest.mark.xslow @given(data=strategies.data()) @pytest.mark.fail_slow(5) @pytest.mark.parametrize( 'func,nfo', [make_xp_pytest_param(nfo.wrapper, nfo) for nfo in _special_funcs]) @pytest.mark.filterwarnings("ignore:invalid value encountered:RuntimeWarning:dask") @pytest.mark.filterwarnings("ignore:divide by zero encountered:RuntimeWarning:dask") @pytest.mark.filterwarnings("ignore:overflow encountered:RuntimeWarning:dask") @pytest.mark.filterwarnings( "ignore:overflow encountered:RuntimeWarning:array_api_strict" ) def test_support_alternative_backends_hypothesis(xp, func, nfo, data): if func.__name__ in {'expn', 'polygamma', 'multigammaln', 'bdtr', 'bdtrc', 'bdtri', 'nbdtr', 'nbdtrc', 'nbdtri', 'pdtri'}: pytest.skip(f"dtypes for {func.__name__} make it a bad fit for this test.") dtype = data.draw(strategies.sampled_from(['float32', 'float64', 'intp'])) positive_only, dtypes = _skip_or_tweak_alternative_backends( xp, nfo, [dtype], (False,)*nfo.n_args ) dtype_np = getattr(np, dtypes[0]) dtype_xp = getattr(xp, dtypes[0]) elements = {'allow_subnormal': False} # Most failures are due to NaN or infinity; uncomment to suppress them # elements['allow_infinity'] = False # elements['allow_nan'] = False if any(positive_only): elements['min_value'] = 0 shapes, _ = data.draw( npst.mutually_broadcastable_shapes(num_shapes=nfo.n_args)) args_np = [data.draw(npst.arrays(dtype_np, shape, elements=elements)) for shape in shapes] args_xp = [xp.asarray(arg, dtype=dtype_xp) for arg in args_np] args_np = [np.asarray(arg, dtype=dtype_np) for arg in args_np] res = nfo.wrapper(*args_xp) # Also wrapped by lazy_xp_function ref = nfo.func(*args_np) # Unwrapped ufunc if ( is_torch(xp) and xpx.default_dtype(xp) == xp.float32 and dtype != "float64" ): # int64 promotes like float32 on torch with default dtype = float32 # cast reference if needed ref = np.float32(ref) # When dtype_np is integer, the output dtype can be float atol = 0 if ref.dtype.kind in 'iu' else 10 * np.finfo(ref.dtype).eps xp_assert_close(res, xp.asarray(ref), atol=atol) @pytest.mark.filterwarnings("ignore:numpy.core is deprecated:DeprecationWarning") @pytest.mark.parametrize("func", [nfo.wrapper for nfo in _special_funcs]) def test_pickle(func): roundtrip = pickle.loads(pickle.dumps(func)) assert roundtrip is func @pytest.mark.parametrize("func", [nfo.wrapper for nfo in _special_funcs]) def test_repr(func): assert func.__name__ in repr(func) assert "locals" not in repr(func) @pytest.mark.skipif( version.parse(np.__version__) < version.parse("2.2"), reason="Can't update ufunc __doc__ when SciPy is compiled vs. NumPy < 2.2") @pytest.mark.parametrize('func', [nfo.wrapper for nfo in _special_funcs]) def test_doc(func): """xp_capabilities updates the docstring in place. Make sure it does so exactly once, including when SCIPY_ARRAY_API is not set. """ match = "has experimental support for Python Array API" assert func.__doc__.count(match) == 1 @pytest.mark.parametrize( 'func,n_args,int_only,is_ufunc', [(nfo.wrapper, nfo.n_args, nfo.int_only, nfo.is_ufunc) for nfo in _special_funcs] ) def test_ufunc_kwargs(func, n_args, int_only, is_ufunc): """Test that numpy-specific out= and dtype= keyword arguments of ufuncs still work when SCIPY_ARRAY_API is set. """ if not is_ufunc: pytest.skip(f"{func.__name__} is not a ufunc.") if int_only is None: int_only = (False, ) * n_args # out= args = [ np.asarray([.1, .2]) if not needs_int else np.asarray([1, 2]) for needs_int in int_only ] out = np.empty(2) y = func(*args, out=out) xp_assert_close(y, out) # out= with out.dtype != args.dtype out = np.empty(2, dtype=np.float32) y = func(*args, out=out) xp_assert_close(y, out) if func.__name__ in {"bdtr", "bdtrc", "bdtri"}: # The below function evaluation will trigger a deprecation warning # with dtype=np.float32. This will go away if the trigger is actually # pulled on the deprecation. return # dtype= y = func(*args, dtype=np.float32) assert y.dtype == np.float32 @make_xp_test_case(special.chdtr) def test_chdtr_gh21311(xp): # the edge case behavior of generic chdtr was not right; see gh-21311 # be sure to test at least these cases # should add `np.nan` into the mix when gh-21317 is resolved x = np.asarray([-np.inf, -1., 0., 1., np.inf]) v = x.reshape(-1, 1) ref = special.chdtr(v, x) res = special.chdtr(xp.asarray(v), xp.asarray(x)) xp_assert_close(res, xp.asarray(ref))