Array API Standard Support: signal#
This page explains some caveats of the signal module and provides (currently
incomplete) tables about the
CPU,
GPU and
JIT support.
Caveats#
JAX and CuPy provide alternative
implementations for some signal functions. When such a function is called, a
decorator decides which implementation to use by inspecting the xp parameter.
Hence, there can be, especially during CI testing, discrepancies in behavior between the default NumPy-based implementation and the JAX and CuPy backends. Skipping the incompatible backends in unit tests, as described in the Adding tests section, is the currently recommended workaround.
The functions are decorated by the code in file
scipy/signal/_support_alternative_backends.py:
1import functools
2import types
3from scipy._lib._array_api import (
4 is_cupy, is_jax, scipy_namespace_for, SCIPY_ARRAY_API, xp_capabilities
5)
6
7from ._signal_api import * # noqa: F403
8from . import _signal_api
9from . import _delegators
10__all__ = _signal_api.__all__
11
12
13MODULE_NAME = 'signal'
14
15# jax.scipy.signal has only partial coverage of scipy.signal, so we keep the list
16# of functions we can delegate to JAX
17# https://jax.readthedocs.io/en/latest/jax.scipy.html
18JAX_SIGNAL_FUNCS = [
19 'fftconvolve', 'convolve', 'convolve2d', 'correlate', 'correlate2d',
20 'csd', 'detrend', 'istft', 'welch'
21]
22
23# some cupyx.scipy.signal functions are incompatible with their scipy counterparts
24CUPY_BLACKLIST = [
25 'abcd_normalize', 'bessel', 'besselap', 'envelope', 'get_window', 'lfilter_zi',
26 'sosfilt_zi', 'remez',
27]
28
29def delegate_xp(delegator, module_name):
30 def inner(func):
31 @functools.wraps(func)
32 def wrapper(*args, **kwds):
33 try:
34 xp = delegator(*args, **kwds)
35 except TypeError:
36 # object arrays
37 if func.__name__ == "tf2ss":
38 import numpy as np
39 xp = np
40 else:
41 raise
42
43 # try delegating to a cupyx/jax namesake
44 if is_cupy(xp) and func.__name__ not in CUPY_BLACKLIST:
45 # https://github.com/cupy/cupy/issues/8336
46 import importlib
47 cupyx_module = importlib.import_module(f"cupyx.scipy.{module_name}")
48 try:
49 cupyx_func = getattr(cupyx_module, func.__name__)
50 except AttributeError:
51 if func.__name__ != "freqz_sos":
52 raise
53 # CuPy < 14 exposes this under the old SciPy name.
54 cupyx_func = cupyx_module.sosfreqz
55 kwds.pop('xp', None)
56 # cupyx functions have no `device` kwarg (cupy#9848); they create on the
57 # current cuda device, so honor a device request via the
58 # cupy.cuda.Device context manager
59 device = kwds.pop('device', None)
60 if device is not None:
61 with device:
62 return cupyx_func(*args, **kwds)
63 return cupyx_func(*args, **kwds)
64 elif is_jax(xp) and func.__name__ in JAX_SIGNAL_FUNCS:
65 spx = scipy_namespace_for(xp)
66 jax_module = getattr(spx, module_name)
67 jax_func = getattr(jax_module, func.__name__)
68 kwds.pop('xp', None)
69 return jax_func(*args, **kwds)
70 else:
71 # the original function
72 return func(*args, **kwds)
73 return wrapper
74 return inner
75
76
77# Although most of these functions currently exist in CuPy and some in JAX,
78# there are no alternative backend tests for any of them in the current
79# test suite. Each will be documented as np_only until tests are added.
80untested = {
81 "argrelextrema",
82 "argrelmax",
83 "argrelmin",
84 "band_stop_obj",
85 "bode",
86 "check_NOLA",
87 "coherence",
88 "csd",
89 "czt",
90 "czt_points",
91 "dbode",
92 "dfreqresp",
93 "dlsim",
94 "dstep",
95 "find_peaks",
96 "find_peaks_cwt",
97 "freqresp",
98 "iirdesign", # There's no reason this shouldn't work. It just needs tests.
99 "istft",
100 "lombscargle",
101 "lsim",
102 "max_len_seq",
103 "peak_prominences",
104 "peak_widths",
105 "periodogram",
106 "place_poles",
107 "sepfir2d",
108 "ss2tf",
109 "ss2zpk",
110 "step",
111 "sweep_poly",
112 "symiirorder1",
113 "symiirorder2",
114 "tf2ss",
115 "unit_impulse",
116 "zoom_fft",
117 "zpk2ss",
118}
119
120
121def get_default_capabilities(func_name, delegator):
122 if delegator is None or func_name in untested:
123 return xp_capabilities(np_only=True)
124 return xp_capabilities()
125
126bilinear_extra_note = \
127 """CuPy does not accept complex inputs.
128
129 """
130
131uses_choose_conv_extra_note = \
132 """CuPy does not support inputs with ``ndim>1`` when ``method="auto"``
133 but does support higher dimensional arrays for ``method="direct"``
134 and ``method="fft"``.
135
136 """
137
138resample_poly_extra_note = \
139 """CuPy only supports ``padtype="constant"``.
140
141 """
142
143upfirdn_extra_note = \
144 """CuPy only supports ``mode="constant"`` and ``cval=0.0``.
145
146 """
147
148xord_extra_note = \
149 """The ``torch`` backend on GPU does not support the case where
150 `wp` and `ws` specify a Bandstop filter.
151
152 """
153
154convolve2d_extra_note = \
155 """The JAX backend only supports ``boundary="fill"`` and ``fillvalue=0``.
156
157 """
158
159zpk2tf_extra_note = \
160 """The CuPy and JAX backends both support only 1d input.
161
162 """
163
164abcd_normalize_extra_note = \
165 """The result dtype when all array inputs are of integer dtype is the
166 backend's current default floating point dtype.
167
168 """
169
170chirp_extra_note = \
171 """CuPy delegates to ``cupyx.scipy.signal.chirp``, which does not support
172 ``complex=True``.
173
174 """
175
176welch_extra_note = \
177 """Support for CuPy and JAX is provided by delegation to
178 ``cupyx.scipy.signal.welch`` and ``jax.scipy.signal.welch``.
179
180 For single-precision input (``float32`` or ``complex64``), JAX returns the sample
181 frequencies in ``float32``, whereas SciPy and CuPy always return them in
182 ``float64``.
183 """
184
185capabilities_overrides = {
186 "abcd_normalize": xp_capabilities(extra_note=abcd_normalize_extra_note),
187 "bessel": xp_capabilities(cpu_only=True, jax_jit=False, allow_dask_compute=True),
188 "bilinear": xp_capabilities(cpu_only=True, exceptions=["cupy"],
189 jax_jit=False, allow_dask_compute=True,
190 reason="Uses np.polynomial.Polynomial",
191 extra_note=bilinear_extra_note),
192 "bilinear_zpk": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
193 jax_jit=False, allow_dask_compute=True),
194 "butter": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
195 allow_dask_compute=True),
196 "buttord": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
197 jax_jit=False, allow_dask_compute=True,
198 extra_note=xord_extra_note),
199 "cheb1ord": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
200 jax_jit=False, allow_dask_compute=True,
201 extra_note=xord_extra_note),
202 "cheb2ord": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
203 jax_jit=False, allow_dask_compute=True,
204 extra_note=xord_extra_note),
205 "cheby1": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
206 allow_dask_compute=True),
207
208 "cheby2": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
209 allow_dask_compute=True),
210 "chirp": xp_capabilities(extra_note=chirp_extra_note),
211 "cont2discrete": xp_capabilities(np_only=True, exceptions=["cupy"]),
212 "convolve": xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"],
213 allow_dask_compute=True,
214 extra_note=uses_choose_conv_extra_note),
215 "convolve2d": xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"],
216 allow_dask_compute=True,
217 extra_note=convolve2d_extra_note),
218 "correlate": xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"],
219 allow_dask_compute=True,
220 extra_note=uses_choose_conv_extra_note),
221 "correlate2d": xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"],
222 allow_dask_compute=True,
223 extra_note=convolve2d_extra_note),
224 "correlation_lags": xp_capabilities(out_of_scope=True),
225 "cspline1d": xp_capabilities(cpu_only=True, exceptions=["cupy"],
226 jax_jit=False, allow_dask_compute=True),
227 "cspline1d_eval": xp_capabilities(cpu_only=True, exceptions=["cupy"],
228 jax_jit=False, allow_dask_compute=True),
229 "cspline2d": xp_capabilities(cpu_only=True, exceptions=["cupy"],
230 jax_jit=False, allow_dask_compute=True),
231 "deconvolve": xp_capabilities(cpu_only=True, exceptions=["cupy"],
232 jax_jit=False, allow_dask_compute=True),
233 "decimate": xp_capabilities(np_only=True, exceptions=["cupy"]),
234 "detrend": xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"],
235 allow_dask_compute=True),
236 "dimpulse": xp_capabilities(np_only=True, exceptions=["cupy"]),
237 "dlti": xp_capabilities(np_only=True,
238 reason="works in CuPy but delegation isn't set up yet"),
239 "ellip": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
240 allow_dask_compute=True,
241 reason="scipy.special.ellipk"),
242 "ellipord": xp_capabilities(cpu_only=True, exceptions=["cupy"],
243 jax_jit=False, allow_dask_compute=True,
244 reason="scipy.special.ellipk"),
245 "filtfilt": xp_capabilities(cpu_only=True, exceptions=["cupy"],
246 allow_dask_compute=True, jax_jit=False),
247 "findfreqs": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
248 jax_jit=False, allow_dask_compute=True),
249 "firls": xp_capabilities(cpu_only=True, allow_dask_compute=True, jax_jit=False,
250 reason="lstsq"),
251 "firwin": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
252 jax_jit=False, allow_dask_compute=True),
253 "firwin2": xp_capabilities(cpu_only=True, exceptions=["cupy"],
254 jax_jit=False, allow_dask_compute=True,
255 reason="firwin2 uses np.interp"),
256 "freqs": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
257 jax_jit=False, allow_dask_compute=True),
258 "freqs_zpk": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
259 jax_jit=False, allow_dask_compute=True),
260 "freqz": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
261 jax_jit=False, allow_dask_compute=True),
262 "freqz_sos": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
263 jax_jit=False, allow_dask_compute=True),
264 "group_delay": xp_capabilities(cpu_only=True, exceptions=["cupy"],
265 jax_jit=False, allow_dask_compute=True),
266 "invres": xp_capabilities(np_only=True, exceptions=["cupy"]),
267 "invresz": xp_capabilities(np_only=True, exceptions=["cupy"]),
268 "iirfilter": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
269 jax_jit=False, allow_dask_compute=True),
270 "kaiser_atten": xp_capabilities(
271 out_of_scope=True, reason="scalars in, scalars out"
272 ),
273 "kaiser_beta": xp_capabilities(out_of_scope=True, reason="scalars in, scalars out"),
274 "kaiserord": xp_capabilities(out_of_scope=True, reason="scalars in, scalars out"),
275 "lfilter": xp_capabilities(cpu_only=True, exceptions=["cupy"],
276 allow_dask_compute=True, jax_jit=False),
277 "lfilter_zi": xp_capabilities(cpu_only=True, allow_dask_compute=True,
278 jax_jit=False),
279 "lfiltic": xp_capabilities(cpu_only=True, exceptions=["cupy"],
280 allow_dask_compute=True, jax_jit=False),
281 "lp2bp": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
282 allow_dask_compute=True, jax_jit=False),
283 "lp2bp_zpk": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
284 allow_dask_compute=True, jax_jit=False),
285 "lp2bs": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
286 allow_dask_compute=True, jax_jit=False),
287 "lp2bs_zpk": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
288 allow_dask_compute=True, jax_jit=False),
289 "lp2lp": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
290 allow_dask_compute=True, jax_jit=False),
291 "lp2lp_zpk": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
292 allow_dask_compute=True, jax_jit=False),
293 "lp2hp": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
294 allow_dask_compute=True, jax_jit=False),
295 "lp2hp_zpk": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
296 allow_dask_compute=True, jax_jit=False),
297 "lti": xp_capabilities(np_only=True,
298 reason="works in CuPy but delegation isn't set up yet"),
299 "medfilt": xp_capabilities(cpu_only=True, exceptions=["cupy"],
300 allow_dask_compute=True, jax_jit=False,
301 reason="uses scipy.ndimage.rank_filter"),
302 "medfilt2d": xp_capabilities(cpu_only=True, exceptions=["cupy"],
303 allow_dask_compute=True, jax_jit=False,
304 reason="c extension module"),
305 "minimum_phase": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
306 allow_dask_compute=True, jax_jit=False),
307 "normalize": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
308 jax_jit=False, allow_dask_compute=True),
309 "oaconvolve": xp_capabilities(
310 cpu_only=True, exceptions=["cupy", "torch"],
311 xfail_backends=[("dask.array", "wrong answer")],
312 ),
313 "order_filter": xp_capabilities(cpu_only=True, exceptions=["cupy"],
314 allow_dask_compute=True, jax_jit=False,
315 reason="uses scipy.ndimage.rank_filter"),
316 "qspline1d": xp_capabilities(cpu_only=True, exceptions=["cupy"],
317 jax_jit=False, allow_dask_compute=True),
318 "qspline1d_eval": xp_capabilities(cpu_only=True, exceptions=["cupy"],
319 jax_jit=False, allow_dask_compute=True),
320 "qspline2d": xp_capabilities(np_only=True, exceptions=["cupy"]),
321 "remez": xp_capabilities(cpu_only=True, allow_dask_compute=True, jax_jit=False),
322 "resample_poly": xp_capabilities(
323 cpu_only=True, exceptions=["cupy"],
324 jax_jit=False, skip_backends=[("dask.array", "XXX something in dask")],
325 extra_note=resample_poly_extra_note,
326 ),
327 "residue": xp_capabilities(np_only=True, exceptions=["cupy"]),
328 "residuez": xp_capabilities(np_only=True, exceptions=["cupy"]),
329 "savgol_filter": xp_capabilities(cpu_only=True, exceptions=["cupy"],
330 jax_jit=False,
331 reason="convolve1d is cpu-only"),
332 "sos2zpk": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
333 allow_dask_compute=True),
334 "sos2tf": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
335 allow_dask_compute=True),
336 "sosfilt": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
337 allow_dask_compute=True),
338 "sosfilt_zi": xp_capabilities(cpu_only=True, allow_dask_compute=True,
339 jax_jit=False),
340 "sosfiltfilt": xp_capabilities(
341 cpu_only=True, exceptions=["cupy"], jax_jit=False,
342 skip_backends=[
343 (
344 "dask.array",
345 "sosfiltfilt directly sets shape attributes on arrays"
346 " which dask doesn't like"
347 ),
348 ],
349 ),
350 "sosfreqz": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
351 jax_jit=False, allow_dask_compute=True),
352 "spline_filter": xp_capabilities(cpu_only=True, exceptions=["cupy"],
353 jax_jit=False, allow_dask_compute=True),
354 "tf2sos": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
355 allow_dask_compute=True),
356 "tf2zpk": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
357 allow_dask_compute=True),
358 "unique_roots": xp_capabilities(np_only=True, exceptions=["cupy"]),
359 "upfirdn": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
360 allow_dask_compute=True,
361 reason="Cython implementation",
362 extra_note=upfirdn_extra_note),
363 "vectorstrength": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
364 allow_dask_compute=True, jax_jit=False),
365 "welch": xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"],
366 allow_dask_compute=True,
367 extra_note=welch_extra_note),
368 "wiener": xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"],
369 allow_dask_compute=True, jax_jit=False,
370 reason="uses scipy.signal.correlate"),
371 "zpk2sos": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
372 allow_dask_compute=True),
373 "zpk2tf": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
374 allow_dask_compute=True,
375 extra_note=zpk2tf_extra_note),
376 "spectrogram": xp_capabilities(out_of_scope=True), # legacy
377 "stft": xp_capabilities(out_of_scope=True), # legacy
378 "istft": xp_capabilities(out_of_scope=True), # legacy
379 "check_COLA": xp_capabilities(out_of_scope=True), # legacy
380}
381
382
383# ### decorate ###
384for obj_name in _signal_api.__all__:
385 bare_obj = getattr(_signal_api, obj_name)
386 delegator = getattr(_delegators, obj_name + "_signature", None)
387
388 if SCIPY_ARRAY_API and delegator is not None:
389 f = delegate_xp(delegator, MODULE_NAME)(bare_obj)
390 else:
391 f = bare_obj
392
393 if not isinstance(f, types.ModuleType):
394 capabilities = capabilities_overrides.get(
395 obj_name, get_default_capabilities(obj_name, delegator)
396 )
397 f = capabilities(f) # pyrefly:ignore[not-callable]
398
399 # add the decorated function to the namespace, to be imported in __init__.py
400 vars()[obj_name] = f
Note that a function will only be decorated if the environment variable
SCIPY_ARRAY_API is set and its signature is listed in the file
scipy/signal/_delegators.py. E.g., for firwin, the signature
function looks like this:
339def firwin_signature(numtaps, cutoff, *args, **kwds):
340 if isinstance(cutoff, int | float):
341 xp = np_compat
342 else:
343 xp = array_namespace(cutoff)
344 return xp
Support on CPU#
Legend
✔️ = supported
✖ = unsupported
N/A = out-of-scope
function/class |
torch |
jax |
dask |
|---|---|---|---|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
N/A |
N/A |
N/A |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
N/A |
N/A |
N/A |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
N/A |
N/A |
N/A |
|
N/A |
N/A |
N/A |
|
N/A |
N/A |
N/A |
|
N/A |
N/A |
N/A |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
N/A |
N/A |
N/A |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
N/A |
N/A |
N/A |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
Support on GPU#
Legend
✔️ = supported
✖ = unsupported
N/A = out-of-scope
function/class |
cupy |
torch |
jax |
|---|---|---|---|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✔️ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✔️ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✖ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
N/A |
N/A |
N/A |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✖ |
✔️ |
|
✔️ |
✖ |
✔️ |
|
✔️ |
✖ |
✔️ |
|
✔️ |
✖ |
✔️ |
|
N/A |
N/A |
N/A |
|
✖ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✖ |
✔️ |
|
✖ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✖ |
|
✔️ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✖ |
|
✔️ |
✔️ |
✖ |
|
✔️ |
✔️ |
✖ |
|
✔️ |
✔️ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
N/A |
N/A |
N/A |
|
N/A |
N/A |
N/A |
|
N/A |
N/A |
N/A |
|
N/A |
N/A |
N/A |
|
✔️ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✖ |
|
✔️ |
✔️ |
✖ |
|
✔️ |
✔️ |
✖ |
|
✔️ |
✔️ |
✖ |
|
✔️ |
✔️ |
✖ |
|
✔️ |
✔️ |
✖ |
|
✔️ |
✔️ |
✖ |
|
✔️ |
✔️ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✔️ |
✖ |
|
✔️ |
✔️ |
✖ |
|
✔️ |
✔️ |
✖ |
|
✔️ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✔️ |
✖ |
|
N/A |
N/A |
N/A |
|
✔️ |
✖ |
✖ |
|
✔️ |
✔️ |
✔️ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
N/A |
N/A |
N/A |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✔️ |
✔️ |
✖ |
|
✔️ |
✖ |
✔️ |
|
✖ |
✖ |
✖ |
|
✔️ |
✖ |
✔️ |
|
✖ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
|
✖ |
✖ |
✖ |
|
✔️ |
✖ |
✖ |
Support with JIT#
Legend
✔️ = supported
✖ = unsupported
N/A = out-of-scope
function/class |
jax |
|---|---|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✔️ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✔️ |
|
✖ |
|
✖ |
|
✖ |
|
✔️ |
|
✖ |
|
✖ |
|
✔️ |
|
✖ |
|
✔️ |
|
✖ |
|
✖ |
|
✖ |
|
N/A |
|
✖ |
|
✔️ |
|
✔️ |
|
✖ |
|
✖ |
|
✖ |
|
✔️ |
|
✔️ |
|
✔️ |
|
✔️ |
|
N/A |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✔️ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✔️ |
|
✖ |
|
✔️ |
|
✔️ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✔️ |
|
✔️ |
|
✔️ |
|
✔️ |
|
✔️ |
|
✖ |
|
✔️ |
|
✔️ |
|
✔️ |
|
✖ |
|
✖ |
|
✔️ |
|
✔️ |
|
✔️ |
|
✖ |
|
✖ |
|
N/A |
|
N/A |
|
N/A |
|
N/A |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✔️ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✔️ |
|
✖ |
|
✖ |
|
✖ |
|
✔️ |
|
✖ |
|
✔️ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
N/A |
|
✖ |
|
✔️ |
|
✖ |
|
✖ |
|
✖ |
|
N/A |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✔️ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |