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 brainunit.math, brainunit.linalg, brainunit.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 brainunit 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 brainunit._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 or not mapped by the automated sweep. 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, cupy, torch, dask, ndonnx. Backends shown but not tested: none.

High-level summary#

Per-subpackage rating#

Subpackage

numpy

jax

cupy

torch

dask

ndonnx

brainunit.math

Mostly ⚠

Full ✓

Mostly ⚠

Partial ⚠

Partial ⚠

Partial ⚠

brainunit.linalg

Mostly ⚠

Full ✓

Mostly ⚠

Mostly ⚠

Partial ⚠

Limited ✗

brainunit.fft

Full ✓

Full ✓

Full ✓

Full ✓

Full ✓

Limited ✗

Quantity methods

Mostly ⚠

Mostly ⚠

Mostly ⚠

Partial ⚠

Partial ⚠

Partial ⚠

brainunit.lax

JAX-only 🅙

Full ✓

JAX-only 🅙

JAX-only 🅙

JAX-only 🅙

JAX-only 🅙

brainunit.autograd

JAX-only 🅙

Full ✓

JAX-only 🅙

JAX-only 🅙

JAX-only 🅙

JAX-only 🅙

brainunit.sparse

JAX-only 🅙

Full ✓

JAX-only 🅙

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 (brainunit.lax, brainunit.autograd, brainunit.sparse) require this backend.

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

  • cupy — measured on a CUDA device through array_api_compat.cupy. CuPy 14.1.1; device NVIDIA GeForce RTX 3060 Laptop GPU; CUDA driver/runtime 13030 / 12090. Cells and summary ratings below are generated from the measured sweep.

  • torch — through array_api_compat.torch. The torch array-API surface lacks several ops brainunit 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 brainunit._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 brainunit._jax_guard.require_jax_backend, which raises BackendError on any non-jax mantissa.

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

Probe result

numpy

jax

cupy

torch

dask

ndonnx

all functions

🅙

🅙

🅙

🅙

🅙

List of brainunit.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

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

Probe result

numpy

jax

cupy

torch

dask

ndonnx

all functions

List of brainunit.autograd functions
  • grad

  • hessian

  • jacfwd

  • jacobian

  • jacrev

  • value_and_grad

  • vector_grad

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

Probe result

numpy

jax

cupy

torch

dask

ndonnx

all functions

🅙

🅙

🅙

🅙

List of brainunit.sparse functions
  • COO

  • CSC

  • CSR

  • SparseMatrix

  • coo_fromdense

  • coo_todense

  • csc_fromdense

  • csc_todense

  • csr_fromdense

  • csr_todense

brainunit.math#

Public callables in brainunit.math that go through the multi-backend dispatcher. Grouped by the same unit semantics used by the math API reference.

Array Creation and Conversion#

brainunit.math — Array Creation and Conversion#

Function

numpy

jax

cupy

torch

dask

ndonnx

full

eye

identity

[1]

[2]

[3]

tri

[4]

[5]

empty

ones

zeros

full_like

[6]

diag

[7]

tril

triu

empty_like

[8]

ones_like

[9]

zeros_like

[10]

fill_diagonal

[11]

[12]

[13]

array

asarray

arange

linspace

[14]

logspace

[15]

[16]

[17]

meshgrid

[18]

vander

[19]

[20]

tril_indices

tril_indices_from

[21]

[22]

triu_indices

triu_indices_from

[23]

[24]

from_numpy

as_numpy

tree_ones_like

[9]

tree_zeros_like

[10]

Unit-preserving Operations#

brainunit.math — Unit-preserving Operations#

Function

numpy

jax

cupy

torch

dask

ndonnx

row_stack

[25]

concatenate

[26]

stack

vstack

[25]

hstack

[27]

dstack

[28]

column_stack

[29]

[30]

block

[31]

[32]

[33]

append

[34]

[35]

split

[36]

[37]

[38]

array_split

[36]

[37]

[38]

dsplit

[39]

[40]

hsplit

[41]

[42]

vsplit

[43]

[44]

atleast_1d

[45]

atleast_2d

[46]

atleast_3d

[47]

broadcast_arrays

broadcast_to

reshape

[48]

moveaxis

transpose

[49]

[50]

swapaxes

[51]

tile

repeat

flip

fliplr

[52]

flipud

[53]

roll

expand_dims

squeeze

[54]

[54]

sort

max

min

amax

amin

diagflat

[55]

[56]

diagonal

[57]

choose

[58]

[59]

[60]

[61]

ravel

[62]

[63]

flatten

unflatten

remove_diag

astype

real

[64]

imag

[64]

conj

[64]

conjugate

[65]

[66]

[64]

negative

positive

abs

sum

nancumsum

[67]

[68]

nansum

[69]

cumsum

[70]

ediff1d

[71]

[72]

absolute

[73]

fabs

[74]

[75]

median

[76]

[77]

[78]

nanmin

[79]

[80]

nanmax

[81]

[82]

ptp

[83]

[84]

[85]

average

[86]

[87]

mean

std

nanmedian

[88]

[89]

[90]

nanmean

[91]

nanstd

[92]

[93]

diff

rot90

[94]

intersect1d

[95]

[96]

[97]

nan_to_num

[98]

percentile

[99]

[100]

nanpercentile

[101]

[102]

[103]

quantile

[104]

[105]

nanquantile

[106]

[107]

[108]

round

[109]

around

[110]

[111]

rint

[112]

[113]

floor

ceil

trunc

fix

modf

[114]

[115]

fmod

[116]

mod

[117]

[118]

copysign

[119]

remainder

maximum

minimum

fmax

[120]

fmin

[121]

lcm

[122]

[123]

gcd

[124]

[125]

trace

[126]

add

subtract

nextafter

[119]

promote_dtypes

interp

[127]

[128]

[129]

clip

histogram

[130]

[131]

compress

[132]

[133]

extract

[134]

[135]

take

[136]

[136]

[136]

[136]

select

[137]

[138]

where

unique

[139]

[140]

gather

[141]

[142]

ldexp

[143]

Unit-changing Operations#

brainunit.math — Unit-changing Operations#

Function

numpy

jax

cupy

torch

dask

ndonnx

reciprocal

prod

product

nancumprod

[144]

[145]

nanprod

[146]

[147]

cumprod

[148]

cumproduct

[148]

var

nanvar

[149]

[150]

cbrt

[151]

[152]

square

sqrt

multiply

divide

power

[153]

[154]

cross

[155]

[156]

true_divide

[157]

floor_divide

float_power

[158]

divmod

[159]

[160]

convolve

[161]

[162]

[163]

trapezoid

?

?

?

?

?

?

dot

[164]

multi_dot

[165]

[166]

[167]

vdot

[168]

vecdot

inner

[169]

[170]

outer

[171]

kron

[172]

[173]

matmul

tensordot

matrix_power

[174]

[175]

correlate

[176]

[177]

[178]

cov

[179]

Dimensionless-input Operations#

brainunit.math — Dimensionless-input Operations#

Function

numpy

jax

cupy

torch

dask

ndonnx

exprel

[180]

🅙

🅙

🅙

🅙

exp

exp2

[181]

expm1

[119]

log

log10

log1p

[119]

log2

arccos

[182]

arccosh

[183]

arcsin

[184]

arcsinh

[185]

arctan

[186]

arctanh

[187]

cos

cosh

sin

sinc

[188]

sinh

tan

tanh

frexp

[189]

hypot

[119]

arctan2

[190]

logaddexp

logaddexp2

[191]

corrcoef

[192]

bitwise_not

[193]

invert

[194]

[195]

bitwise_and

bitwise_or

bitwise_xor

left_shift

[196]

[197]

right_shift

[198]

[199]

Angle and Phase Operations#

brainunit.math — Angle and Phase Operations#

Function

numpy

jax

cupy

torch

dask

ndonnx

deg2rad

[200]

rad2deg

[201]

degrees

[202]

[203]

radians

[204]

[205]

angle

[64]

Unit-removing Operations#

brainunit.math — Unit-removing Operations#

Function

numpy

jax

cupy

torch

dask

ndonnx

heaviside

[206]

[207]

sign

get_promote_dtypes

iscomplexobj

[208]

[209]

[210]

signbit

[119]

equal

not_equal

greater

greater_equal

less

less_equal

array_equal

[211]

[212]

[213]

isclose

[214]

allclose

[215]

all

any

logical_not

logical_and

logical_or

logical_xor

alltrue

sometrue

bincount

[216]

digitize

[217]

[218]

argsort

argmax

argmin

nanargmax

[219]

[220]

nanargmin

[221]

[222]

argwhere

[223]

nonzero

flatnonzero

[224]

[225]

searchsorted

count_nonzero

diag_indices_from

[226]

[227]

[228]

Activation Functions#

brainunit.math — Activation Functions#

Function

numpy

jax

cupy

torch

dask

ndonnx

relu

[180]

🅙

🅙

🅙

🅙

relu6

[180]

🅙

🅙

🅙

🅙

sigmoid

[180]

🅙

🅙

🅙

🅙

softplus

[180]

🅙

🅙

🅙

🅙

sparse_plus

[180]

🅙

🅙

🅙

🅙

sparse_sigmoid

[180]

🅙

🅙

🅙

🅙

soft_sign

[180]

🅙

🅙

🅙

🅙

silu

[180]

🅙

🅙

🅙

🅙

swish

[180]

🅙

🅙

🅙

🅙

log_sigmoid

[180]

🅙

🅙

🅙

🅙

leaky_relu

hard_sigmoid

[180]

🅙

🅙

🅙

🅙

hard_silu

[180]

🅙

🅙

🅙

🅙

hard_swish

[180]

🅙

🅙

🅙

🅙

hard_tanh

[180]

🅙

🅙

🅙

🅙

elu

[180]

🅙

🅙

🅙

🅙

celu

[180]

🅙

🅙

🅙

🅙

selu

[180]

🅙

🅙

🅙

🅙

gelu

[180]

🅙

🅙

🅙

🅙

glu

[180]

🅙

🅙

🅙

🅙

squareplus

[180]

🅙

🅙

🅙

🅙

mish

[180]

🅙

🅙

🅙

🅙

Einstein Operations#

brainunit.math — Einstein Operations#

Function

numpy

jax

cupy

torch

dask

ndonnx

einreduce

einrearrange

[229]

[50]

einrepeat

einshape

einsum

[180]

[230]

[231]

[232]

[233]

Dtypes, Constants, and Utilities#

brainunit.math — Dtypes, Constants, and Utilities#

Function

numpy

jax

cupy

torch

dask

ndonnx

dtype

finfo

iinfo

issubdtype

result_type

ndim

isreal

[234]

isscalar

isfinite

isinf

isnan

shape

size

broadcast_shapes

gradient

[235]

bartlett

[180]

[236]

[237]

[238]

[239]

blackman

[180]

[236]

[237]

[238]

[239]

hamming

[180]

[236]

[237]

[238]

[239]

hanning

[180]

[236]

[237]

[238]

[239]

kaiser

[180]

[236]

[237]

[238]

[239]

Non-dispatched helpers#

These names live under brainunit.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

brainunit.linalg#

brainunit.linalg#

Function

numpy

jax

cupy

torch

dask

ndonnx

cholesky

[119]

cond

[240]

[241]

cross

[155]

[156]

det

[242]

[243]

diagonal

[244]

dot

[164]

eig

[180]

[245]

[246]

[247]

[248]

eigh

[180]

[249]

[250]

[251]

[252]

eigvals

[180]

[245]

[246]

[247]

[248]

eigvalsh

[180]

[249]

[250]

[251]

[252]

inner

[169]

[170]

inv

[253]

kron

[172]

[173]

lstsq

[254]

[255]

matmul

matrix_norm

[256]

matrix_power

[174]

[175]

matrix_rank

[257]

[258]

matrix_transpose

[259]

multi_dot

[165]

[166]

[167]

norm

[260]

outer

[171]

pinv

[261]

[262]

qr

[255]

slogdet

[263]

[255]

solve

[264]

svd

[180]

[245]

[265]

[266]

[248]

svdvals

[180]

[245]

[265]

[266]

[248]

tensordot

tensorinv

[267]

[268]

tensorsolve

[269]

[270]

trace

[271]

vdot

[168]

vecdot

vector_norm

[272]

brainunit.fft#

Routing varies inside brainunit.fft: _fft_change_unit.py calls brainunit._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.

brainunit.fft#

Function

numpy

jax

cupy

torch

dask

ndonnx

fft

[273]

fft2

[274]

fftfreq

[275]

fftn

[276]

fftshift

[277]

ifft

[278]

ifft2

[279]

ifftn

[280]

ifftshift

[281]

irfft

[282]

irfft2

[283]

irfftn

[284]

rfft

[285]

rfft2

[286]

rfftfreq

[275]

rfftn

[287]

Quantity methods#

Methods on brainunit.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 brainunit’s own guard (). .item reports on dask / ndonnx because the underlying array object does not expose .item().

Quantity public methods#

Function

numpy

jax

cupy

torch

dask

ndonnx

all

any

argmax

argmin

argsort

astype

clip

clone

[288]

[289]

[290]

conj

[119]

conjugate

[119]

copy

[288]

[289]

[290]

cpu

[291]

[292]

[293]

[294]

cross

[295]

[296]

[297]

cuda

[298]

[298]

[298]

[298]

[298]

[298]

cumprod

[299]

[299]

[299]

[299]

[299]

[299]

cumsum

[300]

[70]

diagonal

[301]

[302]

dot

[303]

double

expand_as

expand_dims

factorless

fill

[304]

[292]

[304]

[305]

flatten

float

half

[306]

has_same_unit

in_unit

item

[307]

[307]

max

mean

min

nancumprod

[308]

[308]

[308]

[308]

[308]

[308]

nanprod

[146]

[147]

nonzero

outer

[309]

pow

prod

ptp

[83]

[85]

put

[304]

[292]

[310]

[305]

ravel

[63]

repeat

repr_in_unit

reshape

resize

[311]

[312]

[313]

round

scatter_add

[314]

[305]

scatter_div

[315]

[316]

[305]

scatter_max

[317]

[305]

scatter_min

[318]

[305]

scatter_mul

[319]

[320]

[305]

scatter_sub

[314]

[305]

searchsorted

sort

[304]

[292]

[304]

[294]

split

[321]

[312]

[322]

squeeze

[323]

[324]

std

sum

swapaxes

[325]

take

[326]

tile

[327]

[328]

to

to_cupy

[329]

to_dask

[292]

[294]

to_decimal

to_jax

[245]

[330]

to_ndonnx

[331]

[331]

[332]

[333]

to_numpy

to_torch

[334]

[335]

[336]

tolist

[337]

[338]

trace

[339]

[340]

transpose

[49]

[341]

tree_flatten

tree_unflatten

?

?

?

?

?

?

unsqueeze

update_mantissa

[304]

[292]

[304]

[294]

var

view

[342]

[343]

with_unit

Coverage statistic#

Subpackage

Mapped

Non-dispatched

Unmapped

Total

brainunit.math

294

47

1

342

brainunit.linalg

35

0

0

35

brainunit.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 brainunit.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 brainunit.lax / .autograd / .sparse is gated identically.

Footnotes#