import functools import types from scipy._lib._array_api import ( is_cupy, is_jax, scipy_namespace_for, SCIPY_ARRAY_API, xp_capabilities ) from ._signal_api import * # noqa: F403 from . import _signal_api from . import _delegators __all__ = _signal_api.__all__ MODULE_NAME = 'signal' # jax.scipy.signal has only partial coverage of scipy.signal, so we keep the list # of functions we can delegate to JAX # https://jax.readthedocs.io/en/latest/jax.scipy.html JAX_SIGNAL_FUNCS = [ 'fftconvolve', 'convolve', 'convolve2d', 'correlate', 'correlate2d', 'csd', 'detrend', 'istft', 'welch' ] # some cupyx.scipy.signal functions are incompatible with their scipy counterparts CUPY_BLACKLIST = [ 'lfilter_zi', 'sosfilt_zi', 'get_window', 'besselap', 'envelope', 'remez', 'bessel' ] # freqz_sos is a sosfreqz rename, and cupy does not have the new name yet (in v13.x) CUPY_RENAMES = {'freqz_sos': 'sosfreqz'} def delegate_xp(delegator, module_name): def inner(func): @functools.wraps(func) def wrapper(*args, **kwds): try: xp = delegator(*args, **kwds) except TypeError: # object arrays if func.__name__ == "tf2ss": import numpy as np xp = np else: raise # try delegating to a cupyx/jax namesake if is_cupy(xp) and func.__name__ not in CUPY_BLACKLIST: func_name = CUPY_RENAMES.get(func.__name__, func.__name__) # https://github.com/cupy/cupy/issues/8336 import importlib cupyx_module = importlib.import_module(f"cupyx.scipy.{module_name}") cupyx_func = getattr(cupyx_module, func_name) kwds.pop('xp', None) return cupyx_func(*args, **kwds) elif is_jax(xp) and func.__name__ in JAX_SIGNAL_FUNCS: spx = scipy_namespace_for(xp) jax_module = getattr(spx, module_name) jax_func = getattr(jax_module, func.__name__) kwds.pop('xp', None) return jax_func(*args, **kwds) else: # the original function return func(*args, **kwds) return wrapper return inner # Although most of these functions currently exist in CuPy and some in JAX, # there are no alternative backend tests for any of them in the current # test suite. Each will be documented as np_only until tests are added. untested = { "argrelextrema", "argrelmax", "argrelmin", "band_stop_obj", "check_NOLA", "chirp", "coherence", "csd", "czt_points", "dbode", "dfreqresp", "dlsim", "dstep", "find_peaks", "find_peaks_cwt", "findfreqs", "freqresp", "gausspulse", "lombscargle", "lsim", "max_len_seq", "peak_prominences", "peak_widths", "periodogram", "place_pols", "sawtooth", "sepfir2d", "square", "ss2tf", "ss2zpk", "step", "sweep_poly", "symiirorder1", "symiirorder2", "tf2ss", "unit_impulse", "welch", "zoom_fft", "zpk2ss", } def get_default_capabilities(func_name, delegator): if delegator is None or func_name in untested: return xp_capabilities(np_only=True) return xp_capabilities() bilinear_extra_note = \ """CuPy does not accept complex inputs. """ uses_choose_conv_extra_note = \ """CuPy does not support inputs with ``ndim>1`` when ``method="auto"`` but does support higher dimensional arrays for ``method="direct"`` and ``method="fft"``. """ resample_poly_extra_note = \ """CuPy only supports ``padtype="constant"``. """ upfirdn_extra_note = \ """CuPy only supports ``mode="constant"`` and ``cval=0.0``. """ xord_extra_note = \ """The ``torch`` backend on GPU does not support the case where `wp` and `ws` specify a Bandstop filter. """ convolve2d_extra_note = \ """The JAX backend only supports ``boundary="fill"`` and ``fillvalue=0``. """ zpk2tf_extra_note = \ """The CuPy and JAX backends both support only 1d input. """ capabilities_overrides = { "bessel": xp_capabilities(cpu_only=True, jax_jit=False, allow_dask_compute=True), "bilinear": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False, allow_dask_compute=True, reason="Uses np.polynomial.Polynomial", extra_note=bilinear_extra_note), "bilinear_zpk": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"], jax_jit=False, allow_dask_compute=True), "butter": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False, allow_dask_compute=True), "buttord": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"], jax_jit=False, allow_dask_compute=True, extra_note=xord_extra_note), "cheb1ord": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"], jax_jit=False, allow_dask_compute=True, extra_note=xord_extra_note), "cheb2ord": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"], jax_jit=False, allow_dask_compute=True, extra_note=xord_extra_note), "cheby1": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False, allow_dask_compute=True), "cheby2": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False, allow_dask_compute=True), "cont2discrete": xp_capabilities(np_only=True, exceptions=["cupy"]), "convolve": xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"], allow_dask_compute=True, extra_note=uses_choose_conv_extra_note), "convolve2d": xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"], allow_dask_compute=True, extra_note=convolve2d_extra_note), "correlate": xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"], allow_dask_compute=True, extra_note=uses_choose_conv_extra_note), "correlate2d": xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"], allow_dask_compute=True, extra_note=convolve2d_extra_note), "correlation_lags": xp_capabilities(out_of_scope=True), "cspline1d": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False, allow_dask_compute=True), "cspline1d_eval": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False, allow_dask_compute=True), "cspline2d": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False, allow_dask_compute=True), "czt": xp_capabilities(np_only=True, exceptions=["cupy"]), "deconvolve": xp_capabilities(cpu_only=True, exceptions=["cupy"], allow_dask_compute=True, skip_backends=[("jax.numpy", "item assignment")]), "decimate": xp_capabilities(np_only=True, exceptions=["cupy"]), "detrend": xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"], allow_dask_compute=True), "dimpulse": xp_capabilities(np_only=True, exceptions=["cupy"]), "dlti": xp_capabilities(np_only=True, reason="works in CuPy but delegation isn't set up yet"), "ellip": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False, allow_dask_compute=True, reason="scipy.special.ellipk"), "ellipord": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False, allow_dask_compute=True, reason="scipy.special.ellipk"), "firls": xp_capabilities(cpu_only=True, allow_dask_compute=True, jax_jit=False, reason="lstsq"), "firwin": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"], jax_jit=False, allow_dask_compute=True), "firwin2": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False, allow_dask_compute=True, reason="firwin uses np.interp"), "fftconvolve": xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"]), "freqs": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"], jax_jit=False, allow_dask_compute=True), "freqs_zpk": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"], jax_jit=False, allow_dask_compute=True), "freqz": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"], jax_jit=False, allow_dask_compute=True), "freqz_sos": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"], jax_jit=False, allow_dask_compute=True), "group_delay": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False, allow_dask_compute=True), "hilbert": xp_capabilities( cpu_only=True, exceptions=["cupy", "torch"], skip_backends=[("jax.numpy", "item assignment")], ), "hilbert2": xp_capabilities( cpu_only=True, exceptions=["cupy", "torch"], skip_backends=[("jax.numpy", "item assignment")], ), "invres": xp_capabilities(np_only=True, exceptions=["cupy"]), "invresz": xp_capabilities(np_only=True, exceptions=["cupy"]), "iircomb": xp_capabilities(xfail_backends=[("jax.numpy", "inaccurate")]), "iirfilter": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"], jax_jit=False, allow_dask_compute=True), "kaiser_atten": xp_capabilities( out_of_scope=True, reason="scalars in, scalars out" ), "kaiser_beta": xp_capabilities(out_of_scope=True, reason="scalars in, scalars out"), "kaiserord": xp_capabilities(out_of_scope=True, reason="scalars in, scalars out"), "lfilter": xp_capabilities(cpu_only=True, exceptions=["cupy"], allow_dask_compute=True, jax_jit=False), "lfilter_zi": xp_capabilities(cpu_only=True, allow_dask_compute=True, jax_jit=False), "lfiltic": xp_capabilities(cpu_only=True, exceptions=["cupy"], allow_dask_compute=True), "lp2bp": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"], allow_dask_compute=True, skip_backends=[("jax.numpy", "in-place item assignment")]), "lp2bp_zpk": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"], allow_dask_compute=True, jax_jit=False), "lp2bs": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"], allow_dask_compute=True, skip_backends=[("jax.numpy", "in-place item assignment")]), "lp2bs_zpk": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"], allow_dask_compute=True, jax_jit=False), "lp2lp": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"], allow_dask_compute=True, jax_jit=False), "lp2lp_zpk": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"], allow_dask_compute=True, jax_jit=False), "lp2hp": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"], allow_dask_compute=True, skip_backends=[("jax.numpy", "in-place item assignment")]), "lp2hp_zpk": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"], allow_dask_compute=True, jax_jit=False), "lti": xp_capabilities(np_only=True, reason="works in CuPy but delegation isn't set up yet"), "medfilt": xp_capabilities(cpu_only=True, exceptions=["cupy"], allow_dask_compute=True, jax_jit=False, reason="uses scipy.ndimage.rank_filter"), "medfilt2d": xp_capabilities(cpu_only=True, exceptions=["cupy"], allow_dask_compute=True, jax_jit=False, reason="c extension module"), "minimum_phase": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"], allow_dask_compute=True, jax_jit=False), "normalize": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"], jax_jit=False, allow_dask_compute=True), "oaconvolve": xp_capabilities( cpu_only=True, exceptions=["cupy", "torch"], skip_backends=[("jax.numpy", "fails all around")], xfail_backends=[("dask.array", "wrong answer")], ), "order_filter": xp_capabilities(cpu_only=True, exceptions=["cupy"], allow_dask_compute=True, jax_jit=False, reason="uses scipy.ndimage.rank_filter"), "qspline1d": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False, allow_dask_compute=True), "qspline1d_eval": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False, allow_dask_compute=True), "qspline2d": xp_capabilities(np_only=True, exceptions=["cupy"]), "remez": xp_capabilities(cpu_only=True, allow_dask_compute=True, jax_jit=False), "resample": xp_capabilities( cpu_only=True, exceptions=["cupy"], skip_backends=[ ("dask.array", "XXX something in dask"), ("jax.numpy", "XXX: immutable arrays"), ] ), "resample_poly": xp_capabilities( cpu_only=True, exceptions=["cupy"], jax_jit=False, skip_backends=[("dask.array", "XXX something in dask")], extra_note=resample_poly_extra_note, ), "residue": xp_capabilities(np_only=True, exceptions=["cupy"]), "residuez": xp_capabilities(np_only=True, exceptions=["cupy"]), "savgol_filter": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False, reason="convolve1d is cpu-only"), "sepfir2d": xp_capabilities(np_only=True), "sos2zpk": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False, allow_dask_compute=True), "sos2tf": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False, allow_dask_compute=True), "sosfilt": xp_capabilities(cpu_only=True, exceptions=["cupy"], allow_dask_compute=True), "sosfiltfilt": xp_capabilities( cpu_only=True, exceptions=["cupy"], skip_backends=[ ( "dask.array", "sosfiltfilt directly sets shape attributes on arrays" " which dask doesn't like" ), ("torch", "negative strides"), ("jax.numpy", "sosfilt works in-place"), ], ), "sosfreqz": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"], jax_jit=False, allow_dask_compute=True), "spline_filter": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False, allow_dask_compute=True), "tf2sos": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False, allow_dask_compute=True), "tf2zpk": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False, allow_dask_compute=True), "unique_roots": xp_capabilities(np_only=True, exceptions=["cupy"]), "upfirdn": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False, allow_dask_compute=True, reason="Cython implementation", extra_note=upfirdn_extra_note), "vectorstrength": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"], allow_dask_compute=True, jax_jit=False), "wiener": xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"], allow_dask_compute=True, jax_jit=False, reason="uses scipy.signal.correlate"), "zpk2sos": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False, allow_dask_compute=True), "zpk2tf": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False, allow_dask_compute=True, extra_note=zpk2tf_extra_note), "spectrogram": xp_capabilities(out_of_scope=True), # legacy "stft": xp_capabilities(out_of_scope=True), # legacy "istft": xp_capabilities(out_of_scope=True), # legacy "check_COLA": xp_capabilities(out_of_scope=True), # legacy } # ### decorate ### for obj_name in _signal_api.__all__: bare_obj = getattr(_signal_api, obj_name) delegator = getattr(_delegators, obj_name + "_signature", None) if SCIPY_ARRAY_API and delegator is not None: f = delegate_xp(delegator, MODULE_NAME)(bare_obj) else: f = bare_obj if not isinstance(f, types.ModuleType): capabilities = capabilities_overrides.get( obj_name, get_default_capabilities(obj_name, delegator) ) f = capabilities(f) # add the decorated function to the namespace, to be imported in __init__.py vars()[obj_name] = f