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