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 return cupyx_func(*args, **kwds)
57 elif is_jax(xp) and func.__name__ in JAX_SIGNAL_FUNCS:
58 spx = scipy_namespace_for(xp)
59 jax_module = getattr(spx, module_name)
60 jax_func = getattr(jax_module, func.__name__)
61 kwds.pop('xp', None)
62 return jax_func(*args, **kwds)
63 else:
64 # the original function
65 return func(*args, **kwds)
66 return wrapper
67 return inner
68
69
70# Although most of these functions currently exist in CuPy and some in JAX,
71# there are no alternative backend tests for any of them in the current
72# test suite. Each will be documented as np_only until tests are added.
73untested = {
74 "argrelextrema",
75 "argrelmax",
76 "argrelmin",
77 "band_stop_obj",
78 "bode",
79 "check_NOLA",
80 "coherence",
81 "csd",
82 "czt",
83 "czt_points",
84 "dbode",
85 "dfreqresp",
86 "dlsim",
87 "dstep",
88 "find_peaks",
89 "find_peaks_cwt",
90 "freqresp",
91 "iirdesign", # There's no reason this shouldn't work. It just needs tests.
92 "istft",
93 "lombscargle",
94 "lsim",
95 "max_len_seq",
96 "peak_prominences",
97 "peak_widths",
98 "periodogram",
99 "place_poles",
100 "sepfir2d",
101 "ss2tf",
102 "ss2zpk",
103 "step",
104 "sweep_poly",
105 "symiirorder1",
106 "symiirorder2",
107 "tf2ss",
108 "unit_impulse",
109 "zoom_fft",
110 "zpk2ss",
111}
112
113
114def get_default_capabilities(func_name, delegator):
115 if delegator is None or func_name in untested:
116 return xp_capabilities(np_only=True)
117 return xp_capabilities()
118
119bilinear_extra_note = \
120 """CuPy does not accept complex inputs.
121
122 """
123
124uses_choose_conv_extra_note = \
125 """CuPy does not support inputs with ``ndim>1`` when ``method="auto"``
126 but does support higher dimensional arrays for ``method="direct"``
127 and ``method="fft"``.
128
129 """
130
131resample_poly_extra_note = \
132 """CuPy only supports ``padtype="constant"``.
133
134 """
135
136upfirdn_extra_note = \
137 """CuPy only supports ``mode="constant"`` and ``cval=0.0``.
138
139 """
140
141xord_extra_note = \
142 """The ``torch`` backend on GPU does not support the case where
143 `wp` and `ws` specify a Bandstop filter.
144
145 """
146
147convolve2d_extra_note = \
148 """The JAX backend only supports ``boundary="fill"`` and ``fillvalue=0``.
149
150 """
151
152zpk2tf_extra_note = \
153 """The CuPy and JAX backends both support only 1d input.
154
155 """
156
157abcd_normalize_extra_note = \
158 """The result dtype when all array inputs are of integer dtype is the
159 backend's current default floating point dtype.
160
161 """
162
163chirp_extra_note = \
164 """CuPy delegates to ``cupyx.scipy.signal.chirp``, which does not support
165 ``complex=True``.
166
167 """
168
169welch_extra_note = \
170 """Support for CuPy and JAX is provided by delegation to
171 ``cupyx.scipy.signal.welch`` and ``jax.scipy.signal.welch``.
172
173 For single-precision input (``float32`` or ``complex64``), JAX returns the sample
174 frequencies in ``float32``, whereas SciPy and CuPy always return them in
175 ``float64``.
176 """
177
178capabilities_overrides = {
179 "abcd_normalize": xp_capabilities(extra_note=abcd_normalize_extra_note),
180 "bessel": xp_capabilities(cpu_only=True, jax_jit=False, allow_dask_compute=True),
181 "bilinear": xp_capabilities(cpu_only=True, exceptions=["cupy"],
182 jax_jit=False, allow_dask_compute=True,
183 reason="Uses np.polynomial.Polynomial",
184 extra_note=bilinear_extra_note),
185 "bilinear_zpk": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
186 jax_jit=False, allow_dask_compute=True),
187 "butter": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
188 allow_dask_compute=True),
189 "buttord": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
190 jax_jit=False, allow_dask_compute=True,
191 extra_note=xord_extra_note),
192 "cheb1ord": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
193 jax_jit=False, allow_dask_compute=True,
194 extra_note=xord_extra_note),
195 "cheb2ord": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
196 jax_jit=False, allow_dask_compute=True,
197 extra_note=xord_extra_note),
198 "cheby1": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
199 allow_dask_compute=True),
200
201 "cheby2": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
202 allow_dask_compute=True),
203 "chirp": xp_capabilities(extra_note=chirp_extra_note),
204 "cont2discrete": xp_capabilities(np_only=True, exceptions=["cupy"]),
205 "convolve": xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"],
206 allow_dask_compute=True,
207 extra_note=uses_choose_conv_extra_note),
208 "convolve2d": xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"],
209 allow_dask_compute=True,
210 extra_note=convolve2d_extra_note),
211 "correlate": xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"],
212 allow_dask_compute=True,
213 extra_note=uses_choose_conv_extra_note),
214 "correlate2d": xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"],
215 allow_dask_compute=True,
216 extra_note=convolve2d_extra_note),
217 "correlation_lags": xp_capabilities(out_of_scope=True),
218 "cspline1d": xp_capabilities(cpu_only=True, exceptions=["cupy"],
219 jax_jit=False, allow_dask_compute=True),
220 "cspline1d_eval": xp_capabilities(cpu_only=True, exceptions=["cupy"],
221 jax_jit=False, allow_dask_compute=True),
222 "cspline2d": xp_capabilities(cpu_only=True, exceptions=["cupy"],
223 jax_jit=False, allow_dask_compute=True),
224 "deconvolve": xp_capabilities(cpu_only=True, exceptions=["cupy"],
225 jax_jit=False, allow_dask_compute=True),
226 "decimate": xp_capabilities(np_only=True, exceptions=["cupy"]),
227 "detrend": xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"],
228 allow_dask_compute=True),
229 "dimpulse": xp_capabilities(np_only=True, exceptions=["cupy"]),
230 "dlti": xp_capabilities(np_only=True,
231 reason="works in CuPy but delegation isn't set up yet"),
232 "ellip": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
233 allow_dask_compute=True,
234 reason="scipy.special.ellipk"),
235 "ellipord": xp_capabilities(cpu_only=True, exceptions=["cupy"],
236 jax_jit=False, allow_dask_compute=True,
237 reason="scipy.special.ellipk"),
238 "filtfilt": xp_capabilities(cpu_only=True, exceptions=["cupy"],
239 allow_dask_compute=True, jax_jit=False),
240 "findfreqs": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
241 jax_jit=False, allow_dask_compute=True),
242 "firls": xp_capabilities(cpu_only=True, allow_dask_compute=True, jax_jit=False,
243 reason="lstsq"),
244 "firwin": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
245 jax_jit=False, allow_dask_compute=True),
246 "firwin2": xp_capabilities(cpu_only=True, exceptions=["cupy"],
247 jax_jit=False, allow_dask_compute=True,
248 reason="firwin2 uses np.interp"),
249 "freqs": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
250 jax_jit=False, allow_dask_compute=True),
251 "freqs_zpk": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
252 jax_jit=False, allow_dask_compute=True),
253 "freqz": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
254 jax_jit=False, allow_dask_compute=True),
255 "freqz_sos": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
256 jax_jit=False, allow_dask_compute=True),
257 "group_delay": xp_capabilities(cpu_only=True, exceptions=["cupy"],
258 jax_jit=False, allow_dask_compute=True),
259 "invres": xp_capabilities(np_only=True, exceptions=["cupy"]),
260 "invresz": xp_capabilities(np_only=True, exceptions=["cupy"]),
261 "iircomb": xp_capabilities(xfail_backends=[("jax.numpy", "inaccurate")]),
262 "iirfilter": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
263 jax_jit=False, allow_dask_compute=True),
264 "kaiser_atten": xp_capabilities(
265 out_of_scope=True, reason="scalars in, scalars out"
266 ),
267 "kaiser_beta": xp_capabilities(out_of_scope=True, reason="scalars in, scalars out"),
268 "kaiserord": xp_capabilities(out_of_scope=True, reason="scalars in, scalars out"),
269 "lfilter": xp_capabilities(cpu_only=True, exceptions=["cupy"],
270 allow_dask_compute=True, jax_jit=False),
271 "lfilter_zi": xp_capabilities(cpu_only=True, allow_dask_compute=True,
272 jax_jit=False),
273 "lfiltic": xp_capabilities(cpu_only=True, exceptions=["cupy"],
274 allow_dask_compute=True, jax_jit=False),
275 "lp2bp": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
276 allow_dask_compute=True, jax_jit=False),
277 "lp2bp_zpk": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
278 allow_dask_compute=True, jax_jit=False),
279 "lp2bs": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
280 allow_dask_compute=True, jax_jit=False),
281 "lp2bs_zpk": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
282 allow_dask_compute=True, jax_jit=False),
283 "lp2lp": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
284 allow_dask_compute=True, jax_jit=False),
285 "lp2lp_zpk": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
286 allow_dask_compute=True, jax_jit=False),
287 "lp2hp": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
288 allow_dask_compute=True, jax_jit=False),
289 "lp2hp_zpk": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
290 allow_dask_compute=True, jax_jit=False),
291 "lti": xp_capabilities(np_only=True,
292 reason="works in CuPy but delegation isn't set up yet"),
293 "medfilt": xp_capabilities(cpu_only=True, exceptions=["cupy"],
294 allow_dask_compute=True, jax_jit=False,
295 reason="uses scipy.ndimage.rank_filter"),
296 "medfilt2d": xp_capabilities(cpu_only=True, exceptions=["cupy"],
297 allow_dask_compute=True, jax_jit=False,
298 reason="c extension module"),
299 "minimum_phase": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
300 allow_dask_compute=True, jax_jit=False),
301 "normalize": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
302 jax_jit=False, allow_dask_compute=True),
303 "oaconvolve": xp_capabilities(
304 cpu_only=True, exceptions=["cupy", "torch"],
305 xfail_backends=[("dask.array", "wrong answer")],
306 ),
307 "order_filter": xp_capabilities(cpu_only=True, exceptions=["cupy"],
308 allow_dask_compute=True, jax_jit=False,
309 reason="uses scipy.ndimage.rank_filter"),
310 "qspline1d": xp_capabilities(cpu_only=True, exceptions=["cupy"],
311 jax_jit=False, allow_dask_compute=True),
312 "qspline1d_eval": xp_capabilities(cpu_only=True, exceptions=["cupy"],
313 jax_jit=False, allow_dask_compute=True),
314 "qspline2d": xp_capabilities(np_only=True, exceptions=["cupy"]),
315 "remez": xp_capabilities(cpu_only=True, allow_dask_compute=True, jax_jit=False),
316 "resample_poly": xp_capabilities(
317 cpu_only=True, exceptions=["cupy"],
318 jax_jit=False, skip_backends=[("dask.array", "XXX something in dask")],
319 extra_note=resample_poly_extra_note,
320 ),
321 "residue": xp_capabilities(np_only=True, exceptions=["cupy"]),
322 "residuez": xp_capabilities(np_only=True, exceptions=["cupy"]),
323 "savgol_filter": xp_capabilities(cpu_only=True, exceptions=["cupy"],
324 jax_jit=False,
325 reason="convolve1d is cpu-only"),
326 "sos2zpk": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
327 allow_dask_compute=True),
328 "sos2tf": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
329 allow_dask_compute=True),
330 "sosfilt": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
331 allow_dask_compute=True),
332 "sosfilt_zi": xp_capabilities(cpu_only=True, allow_dask_compute=True,
333 jax_jit=False),
334 "sosfiltfilt": xp_capabilities(
335 cpu_only=True, exceptions=["cupy"], jax_jit=False,
336 skip_backends=[
337 (
338 "dask.array",
339 "sosfiltfilt directly sets shape attributes on arrays"
340 " which dask doesn't like"
341 ),
342 ],
343 ),
344 "sosfreqz": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
345 jax_jit=False, allow_dask_compute=True),
346 "spline_filter": xp_capabilities(cpu_only=True, exceptions=["cupy"],
347 jax_jit=False, allow_dask_compute=True),
348 "tf2sos": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
349 allow_dask_compute=True),
350 "tf2zpk": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
351 allow_dask_compute=True),
352 "unique_roots": xp_capabilities(np_only=True, exceptions=["cupy"]),
353 "upfirdn": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
354 allow_dask_compute=True,
355 reason="Cython implementation",
356 extra_note=upfirdn_extra_note),
357 "vectorstrength": xp_capabilities(cpu_only=True, exceptions=["cupy", "torch"],
358 allow_dask_compute=True, jax_jit=False),
359 "welch": xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"],
360 allow_dask_compute=True,
361 extra_note=welch_extra_note),
362 "wiener": xp_capabilities(cpu_only=True, exceptions=["cupy", "jax.numpy"],
363 allow_dask_compute=True, jax_jit=False,
364 reason="uses scipy.signal.correlate"),
365 "zpk2sos": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
366 allow_dask_compute=True),
367 "zpk2tf": xp_capabilities(cpu_only=True, exceptions=["cupy"], jax_jit=False,
368 allow_dask_compute=True,
369 extra_note=zpk2tf_extra_note),
370 "spectrogram": xp_capabilities(out_of_scope=True), # legacy
371 "stft": xp_capabilities(out_of_scope=True), # legacy
372 "istft": xp_capabilities(out_of_scope=True), # legacy
373 "check_COLA": xp_capabilities(out_of_scope=True), # legacy
374}
375
376
377# ### decorate ###
378for obj_name in _signal_api.__all__:
379 bare_obj = getattr(_signal_api, obj_name)
380 delegator = getattr(_delegators, obj_name + "_signature", None)
381
382 if SCIPY_ARRAY_API and delegator is not None:
383 f = delegate_xp(delegator, MODULE_NAME)(bare_obj)
384 else:
385 f = bare_obj
386
387 if not isinstance(f, types.ModuleType):
388 capabilities = capabilities_overrides.get(
389 obj_name, get_default_capabilities(obj_name, delegator)
390 )
391 f = capabilities(f) # pyrefly:ignore[not-callable]
392
393 # add the decorated function to the namespace, to be imported in __init__.py
394 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 |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✔️ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |
|
✖ |