Feature support matrix#

Note

This page is generated from dev/backend_support_data.json, which is produced by dev/backend_support_sweep.py — an automated sweep that invokes every public function in saiunit.math, saiunit.linalg, saiunit.fft, and every public Quantity method across each locally-installed backend and records the outcome. Re-run the sweep and the renderer to refresh this page.

Cell legend#

Glyph

Meaning

Verified: the call returned a value of the expected backend kind.

Skipped: the backend’s array-API surface does not expose the underlying op, or it rejects a keyword saiunit forwards (e.g. JAX-only precision=).

Failed: the call raised an unexpected exception on this backend.

Works with a caveat (e.g. lazy result on dask, expected BackendError for materialization on dask).

🅙

JAX-only by design — gated by saiunit._jax_guard.require_jax_backend. Raises BackendError on any non-jax backend.

Not applicable to backend dispatch (dtype factories, dimension predicates).

?

Not tested in this report. Cupy is always ? because no CUDA backend was available when this sweep ran. The single unmapped Quantity method (tree_unflatten) is also ? because automated invocation requires a hand-crafted aux/children pair.

Sweep environment. Backends invoked: numpy, jax, torch, dask, ndonnx. Backends shown but not tested: cupy.

High-level summary#

Table 1 Per-subpackage rating#

Subpackage

numpy

jax

cupy

torch

dask

ndonnx

saiunit.math

Full ✓

Full ✓

?

Partial ⚠

Mostly ⚠

Partial ⚠

saiunit.linalg

Full ✓

Full ✓

?

Mostly ⚠

Partial ⚠

Limited ✗

saiunit.fft

Full ✓

Full ✓

?

Full ✓

Full ✓

Limited ✗

Quantity methods

Mostly ⚠

Mostly ⚠

?

Partial ⚠

Partial ⚠

Partial ⚠

saiunit.lax

JAX-only 🅙

Full ✓

?

JAX-only 🅙

JAX-only 🅙

JAX-only 🅙

saiunit.autograd

JAX-only 🅙

Full ✓

?

JAX-only 🅙

JAX-only 🅙

JAX-only 🅙

saiunit.sparse

JAX-only 🅙

Full ✓

?

JAX-only 🅙

JAX-only 🅙

JAX-only 🅙

Rating thresholds: Full ≥ 95 % pass and zero fail; Mostly ≥ 80 % pass; Partial ≥ 30 % pass; Limited < 30 % pass; JAX-only = gated by require_jax_backend.

Backend-specific notes#

  • jax — full feature set; default backend. All JAX-only subpackages (saiunit.lax, saiunit.autograd, saiunit.sparse) require this backend.

  • numpy — eager CPU computation through array_api_compat.numpy. A handful of reductions (amax, amin, mean, nan* variants) fail when saiunit forwards a where=None kwarg numpy can’t interpret. These are listed with footnotes in the math tables below.

  • cupynot tested in this report (no CUDA toolkit in the sweep environment). Cells are ?. Cupy’s array-API surface tracks numpy closely, so support is expected to mirror the numpy column, but this document does not claim it.

  • torch — through array_api_compat.torch. The torch array-API surface lacks several ops saiunit dispatches to (cbrt, digamma, some einops reductions, axes= for n-D FFTs) and rejects JAX-flavored kwargs (precision, symmetrize_input, tol). Affected calls are recorded as skip rather than fail.

  • dask — lazy arrays. Reductions and most array ops succeed but the result remains lazy until .compute(). Per saiunit._base_quantity, the Python casts float(q) / int(q) / operator.index(q) / np.asarray(q) / hash(q) and the Quantity.tolist method raise BackendError to avoid silent materialization — these cells are with the BackendError text in the footnote. Quantity.item on dask raises a different error (the dask Array has no .item() method) so it appears as rather than . Methods like Quantity.float / .double are .astype in disguise and stay lazy on dask, so they pass.

  • ndonnx — symbolic graph-building backend. Many array-API ops (fft.*, several linalg.*, complex / specialty math) are not implemented and surface as skip rows. Saiunit does not encode unit information into the ONNX graph.

JAX-only subpackages#

These subpackages dispatch directly to JAX primitives that have no array-API equivalent. Each entry point is wrapped with saiunit._jax_guard.require_jax_backend, which raises BackendError on any non-jax mantissa.

saiunit.lax — 101 public callable(s); all require jax.

Probe result

numpy

jax

cupy

torch

dask

ndonnx

all functions

🅙

?

🅙

🅙

🅙

List of saiunit.lax functions
  • acos

  • acosh

  • approx_max_k

  • approx_min_k

  • asin

  • asinh

  • atan

  • atan2

  • atanh

  • batch_matmul

  • bessel_i0e

  • bessel_i1e

  • betainc

  • bitcast_convert_type

  • broadcast

  • broadcast_in_dim

  • broadcast_shapes

  • broadcast_to_rank

  • broadcasted_iota

  • cholesky

  • clamp

  • clz

  • collapse

  • complex

  • conv

  • conv_transpose

  • convert_element_type

  • cumlogsumexp

  • cummax

  • cummin

  • cumsum

  • digamma

  • div

  • dot_general

  • dynamic_index_in_dim

  • dynamic_slice

  • dynamic_slice_ind_dim

  • dynamic_update_index_in_dim

  • dynamic_update_slice

  • dynamic_update_slice_in_dim

  • eig

  • eigh

  • eq

  • erf

  • erf_inv

  • erfc

  • fft

  • gather

  • ge

  • gt

  • hessenberg

  • householder_product

  • igamma

  • igamma_grad_a

  • igammac

  • index_in_dim

  • index_take

  • integer_pow

  • iota

  • le

  • lgamma

  • logistic

  • lt

  • lu

  • mul

  • ne

  • neg

  • pad

  • polygamma

  • population_count

  • pow

  • qdwh

  • qr

  • random_gamma_grad

  • reduce

  • reduce_precision

  • rem

  • rsqrt

  • scatter

  • scatter_add

  • scatter_apply

  • scatter_max

  • scatter_min

  • scatter_mul

  • scatter_sub

  • schur

  • shift_left

  • shift_right_arithmetic

  • shift_right_logical

  • slice

  • slice_in_dim

  • sort

  • sort_key_val

  • sub

  • svd

  • top_k

  • triangular_solve

  • tridiagonal

  • tridiagonal_solve

  • zeros_like_array

  • zeta

saiunit.autograd — 7 public callable(s); all require jax.

Probe result

numpy

jax

cupy

torch

dask

ndonnx

all functions

?

List of saiunit.autograd functions
  • grad

  • hessian

  • jacfwd

  • jacobian

  • jacrev

  • value_and_grad

  • vector_grad

saiunit.sparse — 10 public callable(s); all require jax.

Probe result

numpy

jax

cupy

torch

dask

ndonnx

all functions

🅙

?

🅙

🅙

🅙

List of saiunit.sparse functions
  • COO

  • CSC

  • CSR

  • SparseMatrix

  • coo_fromdense

  • coo_todense

  • csc_fromdense

  • csc_todense

  • csr_fromdense

  • csr_todense

saiunit.math#

Public callables in saiunit.math that go through the multi-backend dispatcher. Grouped by source submodule for readability.

array_creation — Array creation#

Table 2 saiunit.math — Array creation#

Function

numpy

jax

cupy

torch

dask

ndonnx

tril_indices

?

triu_indices

?

keep_unit — Unit-preserving#

Table 3 saiunit.math — Unit-preserving#

Function

numpy

jax

cupy

torch

dask

ndonnx

astype

?

change_unit — Unit-changing#

saiunit.math — Unit-changing: no functions in this group.

accept_unitless — Dimensionless-only#

saiunit.math — Dimensionless-only: no functions in this group.

remove_unit — Unit-removing#

saiunit.math — Unit-removing: no functions in this group.

saiunit.math#

Table 4 saiunit.math — saiunit.math#

Function

numpy

jax

cupy

torch

dask

ndonnx

abs

?

absolute

?

[1]

add

?

all

?

allclose

?

[2]

alltrue

?

amax

?

amin

?

angle

?

[3]

any

?

append

?

[4]

[5]

arange

?

arccos

?

[6]

arccosh

?

[7]

arcsin

?

[8]

arcsinh

?

[9]

arctan

?

[10]

arctan2

?

[11]

arctanh

?

[12]

argmax

?

argmin

?

argsort

?

argwhere

?

[13]

around

?

[14]

[15]

array

?

array_equal

?

[16]

[17]

[18]

array_split

?

[19]

[20]

as_numpy

?

asarray

?

atleast_1d

?

[21]

atleast_2d

?

[22]

atleast_3d

?

[23]

average

?

[24]

[25]

bincount

?

[26]

bitwise_and

?

bitwise_not

?

[27]

bitwise_or

?

bitwise_xor

?

block

?

[28]

[29]

broadcast_arrays

?

broadcast_shapes

?

broadcast_to

?

cbrt

?

[30]

[31]

ceil

?

celu

?

🅙

🅙

🅙

choose

?

[32]

[33]

clip

?

column_stack

?

[34]

[35]

compress

?

[36]

[37]

concatenate

?

[38]

conj

?

[3]

conjugate

?

[39]

[40]

[3]

convolve

?

[41]

[42]

[43]

copysign

?

[44]

corrcoef

?

[45]

correlate

?

[46]

[47]

[48]

cos

?

cosh

?

count_nonzero

?

cov

?

[49]

cross

?

[50]

[51]

cumprod

?

[52]

cumproduct

?

[52]

cumsum

?

[53]

deg2rad

?

[54]

degrees

?

[55]

[56]

diag

?

[57]

diag_indices_from

?

[58]

[59]

[60]

diagflat

?

[61]

[62]

diagonal

?

[63]

diff

?

digitize

?

[64]

[65]

divide

?

divmod

?

[66]

[67]

dot

?

[68]

dsplit

?

[69]

[70]

dstack

?

[71]

ediff1d

?

[72]

[73]

einsum

[74]

?

[75]

[74]

[76]

elu

?

🅙

🅙

🅙

empty

?

empty_like

?

equal

?

exp

?

exp2

?

[77]

expand_dims

?

expm1

?

[44]

exprel

?

🅙

🅙

🅙

extract

?

[78]

[79]

eye

?

fabs

?

[80]

[81]

fill_diagonal

?

[82]

[83]

[84]

finfo

?

fix

?

flatnonzero

?

[85]

[86]

flatten

?

flip

?

fliplr

?

[87]

flipud

?

[88]

float_power

?

[89]

floor

?

floor_divide

?

fmax

?

[90]

fmin

?

[91]

fmod

?

[92]

frexp

?

[93]

from_numpy

?

full

?

full_like

?

gather

?

[94]

[95]

gcd

?

[96]

[97]

gelu

?

🅙

🅙

🅙

get_promote_dtypes

?

glu

?

🅙

🅙

🅙

gradient

?

[98]

greater

?

greater_equal

?

hard_sigmoid

?

🅙

🅙

🅙

hard_silu

?

🅙

🅙

🅙

hard_swish

?

🅙

🅙

🅙

hard_tanh

?

🅙

🅙

🅙

heaviside

?

[99]

[100]

histogram

?

[101]

hsplit

?

[102]

[103]

hstack

?

[104]

hypot

?

[44]

identity

?

[105]

[106]

[107]

iinfo

?

imag

?

[3]

inner

?

[108]

[109]

interp

?

[110]

[111]

[112]

intersect1d

?

[113]

[114]

[115]

invert

?

[116]

[117]

isclose

?

[118]

iscomplexobj

?

[119]

[120]

[121]

isfinite

?

isinf

?

isnan

?

isreal

?

[122]

isscalar

?

issubdtype

?

kron

?

[123]

[124]

lcm

?

[125]

[126]

ldexp

?

[127]

leaky_relu

?

left_shift

?

[128]

[129]

less

?

less_equal

?

linspace

?

log

?

log10

?

log1p

?

[44]

log2

?

log_sigmoid

?

🅙

🅙

🅙

logaddexp

?

logaddexp2

?

[130]

logical_and

?

logical_not

?

logical_or

?

logical_xor

?

logspace

?

[131]

[132]

[133]

matmul

?

matrix_power

?

[134]

[135]

max

?

maximum

?

mean

?

median

?

[136]

[137]

meshgrid

?

[138]

min

?

minimum

?

mish

?

🅙

🅙

🅙

mod

?

[139]

[140]

modf

?

[141]

[142]

moveaxis

?

multi_dot

?

[143]

[144]

multiply

?

nan_to_num

?

[145]

nanargmax

?

[146]

[147]

nanargmin

?

[148]

[149]

nancumprod

?

[150]

[151]

nancumsum

?

[152]

[153]

nanmax

?

[154]

[155]

nanmean

?

[156]

nanmedian

?

[157]

[158]

nanmin

?

[159]

[160]

nanpercentile

?

[161]

[162]

nanprod

?

[163]

[164]

nanquantile

?

[165]

[166]

nanstd

?

[167]

[168]

nansum

?

[169]

nanvar

?

[170]

[171]

ndim

?

negative

?

nextafter

?

[44]

nonzero

?

not_equal

?

ones

?

ones_like

?

outer

?

[172]

percentile

?

[173]

[174]

positive

?

power

?

[175]

[176]

prod

?

product

?

promote_dtypes

?

ptp

?

[177]

[178]

quantile

?

[179]

[180]

rad2deg

?

[181]

radians

?

[182]

[183]

ravel

?

[184]

real

?

[3]

reciprocal

?

relu

?

🅙

🅙

🅙

relu6

?

🅙

🅙

🅙

remainder

?

remove_diag

?

[185]

repeat

?

reshape

?

result_type

?

right_shift

?

[186]

[187]

rint

?

[188]

[189]

roll

?

rot90

?

[190]

round

?

row_stack

?

[191]

searchsorted

?

select

?

[192]

[193]

selu

?

🅙

🅙

🅙

shape

?

sigmoid

?

🅙

🅙

🅙

sign

?

signbit

?

[44]

silu

?

🅙

🅙

🅙

sin

?

sinc

?

[194]

sinh

?

size

?

soft_sign

?

🅙

🅙

🅙

softplus

?

🅙

🅙

🅙

sometrue

?

sort

?

sparse_plus

?

🅙

🅙

🅙

sparse_sigmoid

?

🅙

🅙

🅙

split

?

[19]

[20]

sqrt

?

square

?

squareplus

?

🅙

🅙

🅙

squeeze

?

[195]

[195]

stack

?

std

?

subtract

?

sum

?

swapaxes

?

[196]

swish

?

🅙

🅙

🅙

take

?

tan

?

tanh

?

tensordot

?

tile

?

trace

?

[197]

transpose

?

[198]

[199]

tree_ones_like

?

tree_zeros_like

?

tri

?

[200]

[201]

tril

?

tril_indices_from

?

[202]

[203]

triu

?

triu_indices_from

?

[204]

[205]

true_divide

?

[206]

trunc

?

unflatten

?

unique

?

[207]

vander

?

[208]

[209]

var

?

vdot

?

[210]

vecdot

?

vsplit

?

[211]

[212]

vstack

?

[191]

where

?

zeros

?

zeros_like

?

jax.numpy#

Table 5 saiunit.math — jax.numpy#

Function

numpy

jax

cupy

torch

dask

ndonnx

bartlett

?

blackman

?

hamming

?

hanning

?

kaiser

?

numpy#

Table 6 saiunit.math — numpy#

Function

numpy

jax

cupy

torch

dask

ndonnx

dtype

?

brainstate.math#

Table 7 saiunit.math — brainstate.math#

Function

numpy

jax

cupy

torch

dask

ndonnx

einrearrange

?

[213]

[199]

einreduce

?

einrepeat

?

einshape

?

Non-dispatched helpers#

These names live under saiunit.math for convenience but do not dispatch on backend — they are dtype factories (re-exported from jax.numpy) or pure-Python predicates / introspection helpers over Quantity / Unit objects. Behavior is identical on every backend.

  • assert_quantity

  • bfloat16

  • bool_

  • cdouble

  • check_dims

  • check_units

  • complex128

  • complex64

  • complex_

  • csingle

  • display_in_unit

  • double

  • fail_for_dimension_mismatch

  • fail_for_unit_mismatch

  • float16

  • float32

  • float64

  • float_

  • get_dim

  • get_dtype

  • get_magnitude

  • get_mantissa

  • get_or_create_dimension

  • get_unit

  • inexact

  • int16

  • int2

  • int32

  • int4

  • int64

  • int8

  • int_

  • is_dimensionless

  • is_float

  • is_int

  • is_quantity

  • is_unitless

  • maybe_decimal

  • set_exprel_order

  • single

  • uint

  • uint16

  • uint2

  • uint32

  • uint4

  • uint64

  • uint8

saiunit.linalg#

Table 8 saiunit.linalg#

Function

numpy

jax

cupy

torch

dask

ndonnx

cholesky

?

[214]

cond

?

[215]

[216]

cross

?

[50]

[51]

det

?

[217]

[218]

diagonal

?

[63]

dot

?

[68]

eig

?

[219]

[220]

[221]

eigh

?

[222]

[223]

[224]

eigvals

?

[219]

[220]

[221]

eigvalsh

?

[222]

[223]

[224]

inner

?

[108]

[109]

inv

?

[225]

kron

?

[123]

[124]

lstsq

?

[226]

[227]

matmul

?

matrix_norm

?

[228]

matrix_power

?

[134]

[135]

matrix_rank

?

[229]

[230]

matrix_transpose

?

[231]

multi_dot

?

[143]

[144]

norm

?

[232]

outer

?

[172]

pinv

?

[233]

[234]

qr

?

[227]

slogdet

?

[235]

[227]

solve

?

[236]

svd

?

[237]

[238]

[239]

svdvals

?

[237]

[238]

[239]

tensordot

?

tensorinv

?

[240]

[241]

tensorsolve

?

[242]

[243]

trace

?

[197]

vdot

?

[210]

vecdot

?

vector_norm

?

[244]

saiunit.fft#

Routing varies inside saiunit.fft: _fft_change_unit.py calls saiunit._backend.get_backend() directly (e.g. for fftfreq / rfftfreq), while _fft_keep_unit.py delegates to the math package’s _fun_keep_unit_unary helper and inherits its dispatch.

Table 9 saiunit.fft#

Function

numpy

jax

cupy

torch

dask

ndonnx

fft

?

[245]

fft2

?

[246]

fftfreq

?

[247]

fftn

?

[248]

fftshift

?

[249]

ifft

?

[250]

ifft2

?

[251]

ifftn

?

[252]

ifftshift

?

[253]

irfft

?

[254]

irfft2

?

[255]

irfftn

?

[256]

rfft

?

[257]

rfft2

?

[258]

rfftfreq

?

[247]

rfftn

?

[259]

Quantity methods#

Methods on saiunit.Quantity itself. .to_<backend>() methods ignore the current backend and convert to the named one — cells show when the target backend isn’t installed in the sweep environment.

Materialization is documented above (see Backend-specific notes). Quantity.tolist on dask is the one method that raises BackendError from saiunit’s own guard (). .item reports on dask / ndonnx because the underlying array object does not expose .item().

Table 10 Quantity public methods#

Function

numpy

jax

cupy

torch

dask

ndonnx

all

?

any

?

argmax

?

argmin

?

argsort

?

astype

?

clip

?

clone

?

[260]

[261]

[262]

conj

?

[44]

conjugate

?

[44]

copy

?

[260]

[261]

[262]

cpu

?

[263]

[264]

[265]

cross

?

[266]

[267]

[268]

cuda

[269]

[269]

?

[269]

[269]

[269]

cumprod

[270]

[270]

?

[270]

[270]

[270]

cumsum

?

[271]

[53]

diagonal

?

[272]

[273]

dot

?

[274]

double

?

expand_as

?

expand_dims

?

factorless

?

fill

?

[263]

[275]

[276]

flatten

?

float

?

half

?

[277]

has_same_unit

?

in_unit

?

item

?

[278]

[278]

max

?

mean

?

min

?

nancumprod

[279]

[279]

?

[279]

[279]

[279]

nanprod

?

[163]

[164]

nonzero

?

outer

?

[280]

pow

?

prod

?

ptp

?

[177]

[178]

put

?

[263]

[281]

[276]

ravel

?

[184]

repeat

?

repr_in_unit

?

reshape

?

resize

?

[263]

[275]

[282]

round

?

scatter_add

?

[283]

[276]

scatter_div

?

[284]

[276]

scatter_max

?

[285]

[276]

scatter_min

?

[286]

[276]

scatter_mul

?

[287]

[276]

scatter_sub

?

[283]

[276]

searchsorted

?

sort

?

[263]

[275]

[265]

split

?

[288]

[289]

squeeze

?

[290]

[291]

std

?

sum

?

swapaxes

?

[292]

take

?

[293]

tile

?

[294]

[295]

to

?

to_cupy

[296]

[296]

?

[296]

[296]

[296]

to_dask

?

[263]

[265]

to_decimal

?

to_jax

?

[297]

to_ndonnx

[298]

?

[299]

[300]

to_numpy

?

to_torch

?

[301]

[302]

tolist

?

[303]

[304]

trace

?

[305]

[306]

transpose

?

[198]

[307]

tree_flatten

?

tree_unflatten

?

?

?

?

?

?

unsqueeze

?

update_mantissa

?

[263]

[275]

[265]

var

?

view

?

[308]

[309]

with_unit

?

Coverage statistic#

Subpackage

Mapped

Non-dispatched

Unmapped

Total

saiunit.math

294

47

0

341

saiunit.linalg

35

0

0

35

saiunit.fft

16

0

0

16

Quantity

78

0

1

79

Mapped = functions the sweep actually invoked. Non-dispatched = type factories / predicates that don’t go through backend dispatch. Unmapped = no call pattern registered (will appear as ? in tables).

How this was produced#

dev/backend_support_sweep.py walks every public callable in the subpackages above, picks a calling pattern from an in-script registry, and invokes the function under with saiunit.using_backend(b) for each backend b in the local environment. Outcomes are classified as pass / skip / fail / warn / unmapped / na and written to dev/backend_support_data.json.

dev/backend_support_render.py (this script’s source) reads that JSON and emits the rst file you are currently reading. To refresh:

PYTHONPATH=. python dev/backend_support_sweep.py
PYTHONPATH=. python dev/backend_support_render.py

JAX-only subpackages are probed with one representative function per subpackage rather than enumerated — every entry point in saiunit.lax / .autograd / .sparse is gated identically.

Footnotes#