Release Notes#
Version 0.3.1#
Highlights#
Version 0.3.1 is a compatibility release that makes saiunit.Quantity
work with Matplotlib out of the box. Plotting a Quantity no longer
requires any setup — the unit converter is registered automatically when
saiunit is imported and Matplotlib is installed, so ax.plot,
ax.scatter, axis limits, reference lines, histograms, error bars,
box/violin plots, stack plots, hex bins, and pie charts accept
quantities directly and label axes with their units.
New features#
Quantity-aware
errorbar,boxplot,violinplot,stackplot,hexbin, andpie. These Axes methods coerce their inputs to arrays before Matplotlib’s unit converter runs, so a bareQuantitypreviously raised.enable_matplotlib_support()now wraps each one: Quantity arguments are converted to plain magnitudes via the public API and the relevant axis is labeled with the unit, while calls without anyQuantitypass through untouched. Error-bar magnitudes (xerr/yerr) and stacked series are rescaled into the axis unit;piewedge sizes andhexbinCcolour values, which have no axis-unit meaning, are reduced to their magnitudes. If a future Matplotlib release changes one of these method signatures, the affected wrapper raises a clear error naming the version the first time it is called with aQuantity(other functions andimport saiunitare unaffected).
Bug fixes#
Matplotlib support is now enabled automatically on import. Earlier releases registered the
Quantityconverter only when the caller invokedenable_matplotlib_support()explicitly; the documentedexamples/matplotlib_quantity_basics.pyomitted that call and raisedTypeError: Only dimensionless quantities can be converted to NumPy arrays. The converter now registers at import time. Callingenable_matplotlib_support()remains supported for re-registering after the registry is cleared or for targeting a custom units module.Scalar (0-d) quantities work as axis limits and reference lines.
Quantity.__iter__now raises eagerly on 0-d inputs, sonumpy.iterablereportsFalseexactly as it does for a 0-dndarray. This unblocksax.set_xlim/ax.set_ylim,ax.axhline, andax.axvlinewith scalar quantities, which previously crashed inside Matplotlib’s_is_natively_supportedcheck.ax.histaccepts quantities and keeps their units. Histogram and other functions coerce their data withmatplotlib.cbook._reshape_2Dbefore the unit converter runs.enable_matplotlib_support()now wraps that helper soQuantitydata (single arrays or lists of per-dataset arrays) flows through to unit processing instead of failing thenumpy.asanyarraycoercion, and the histogram axis is labeled with the quantity’s unit.ax.axhspan/ax.axvspanaccept quantity bounds. The converter now passes plain (non-Quantity) values through unchanged, fixing the double-conversionConversionErrorraised when Matplotlib re-converts already-scaled patch coordinates. This also lets callers mix bare numbers onto a unit-bearing axis, matching Matplotlib’s own semantics.
Known limitations#
Two-dimensional scalar-field functions (
imshow,contourf,pcolormesh,quiver) treat aQuantityas colour/magnitude data rather than an axis coordinate, and Matplotlib offers no unit hook there. Pass.to_decimal(unit)and label the colour bar manually.
Tests#
Expanded the Matplotlib compatibility suite to 43 tests covering auto-registration,
plot/scatteraxis-unit inference, scalar axis limits and reference lines,axhspan/axvspanand plain-number passthrough, cross-unit conversion and incompatible-unit errors, JAX-backed quantities,histunit preservation, theerrorbar/boxplot/violinplot/stackplot/hexbin/piewrappers, non-Quantity identity, and the fail-loud signature guard.Closed test-coverage gaps across the public surface. Added colocated test files for previously untested modules and functions:
linalgsvdvals/eigvalsunit propagation;mathfrexp,promote_dtypes,iscomplexobj, thealltrue/sometrue/cumproductaliases, andeinshapeparsing; thejax_onlydecorator guard;is_unit_equal; 15 previously-untested physical constants; every unit shortcut and_aliasre-export; and smoke tests for the_jax_compat,_compatible_import, and_typingshims.
Version 0.3.0#
Highlights#
Version 0.3.0 completes the multi-backend transition started in 0.2.2. The
entire public surface of saiunit.math, saiunit.linalg, saiunit.fft,
and Quantity (including Quantity.at and every scatter operation) now
dispatches through the array-API standard, so the same source code runs
unchanged on numpy, jax, cupy, torch, dask, and
ndonnx. A centralized argument-adaptation layer handles the
per-backend signature differences that previously produced TypeError
on calls that the underlying operation supported.
This release also lands a static-typing pass — mypy saiunit/ reports
zero errors across 111 source files, the new type_check CI job blocks
regressions, and the package ships a PEP 561 py.typed marker so
downstream type checkers honor the inline annotations.
Breaking changes#
Removed the
compatible_with_equinoxtoggle and thecompat_with_equinoxglobal. The Equinox compatibility shim, the Equinox branch inQuantity.__pow__, and the associated tests, docs, and mypy overrides have been deleted. Equinox now works against saiunit via the same path as every other consumer.Removed the
promote_integerskeyword fromsaiunit.math.sumandsaiunit.math.prod(and the**kwargscatch-all that smuggled it through). This argument is JAX-only; forwarding it raisedTypeErroron every non-JAX backend. Callers that need it should calljax.numpy.sum/jax.numpy.proddirectly.Quantity(np.ndarray)no longer silently lifts to JAX in code paths that previously did so.Quantity(mantissa, dtype=...), theQuantity.mantissasetter, andQuantity([q1, q2, ...])all now honor the active backend (or the input’s existing backend) instead of defaulting to JAX whenever JAX is installed.Quantity.strides/.flat/.T/.mTraiseBackendErroron lazy mantissas (dask,ndonnx) instead of silently calling.compute()/.unwrap_numpy().Quantity.atonndonnxmantissas raisesBackendError. ndonnx builds a symbolic ONNX graph and has no notion of a functional in-place update; call.to_numpy()first to materialize.
New features#
Multi-backend Quantity.at and scatter dispatch#
All nine
.at[idx].<op>(...)operations —get,set,add,multiply/mul,divide/div,power,min,max,apply— now work onnumpy,jax,cupy,torch, anddask. Unit-tracking semantics are unchanged: the same dimension/magnitude checks fire regardless of backend, and the result preserves the original mantissa backend.Unified scatter dispatch (
saiunit._scatter).Quantity.atandQuantity.__setitem__/scatter_add/scatter_mul/scatter_div/scatter_max/scatter_minall route through one backend-aware module.CustomArray.__setitem__shares the same path, so user subclasses with non-JAX mantissas work without a JAX install.Emulated
modeandfill_valueon non-JAX backends. JAX’smode('promise_in_bounds'/'clip'/'drop'/'fill') andfill_valuearguments are emulated onnumpy/cupy/torchfor scalar-int and 1D-integer-array indices, sox.at[20].get(mode='fill', fill_value=-1 * u.mV)runs unchanged across backends.
Complete multi-backend coverage for saiunit.math / linalg / fft#
Every direct
jnp.*/jax.tree.*/jax.lax.*call insaiunit/math/_fun_{change,keep,remove,array_creation,accept}_unit.py,saiunit/_misc.py,saiunit/fft/, andsaiunit/linalg/now goes throughget_backend()+_resolve_op()so the same code path works on numpy/torch/dask/ndonnx in addition to JAX.Mechanical replacements where backends diverge:
remove_diaguses~xp.eyeinstead ofjnp.fill_diagonal(which numpy implements in-place returningNone);meshgridusesxp.broadcast_to+reshapeinstead ofjax.lax.broadcast_in_dim;compress/extract/unique/where(cond)/searchsorteddrop JAX-only kwargs (size,fill_value,method) when the backend doesn’t accept them.
Centralized dispatch layer#
New
_dispatch_callpipeline insaiunit.math._fun_keep_unitcombines two adaptation layers:_filter_unsupported_kwargs(signature introspection) and_safe_call(TypeError-message parsing for C-bindings such as torch). All dispatch helpers (_fun_keep_unit_*,_fun_change_unit_*,_fun_remove_unit_*,_fun_logic_*,_fun_accept_unitless_*,_fun_unitless_binary) route through it, and the_fun_array_creationcall sites use a matching_safe_call_xpshim.Centralized
None-kwarg stripping (_strip_none_kwargs) in the same dispatch helpers — wrapper call sites pass kwargs naturally instead of inline_drop_none(...)plumbing.Per-wrapper signature fixes so the generated backend-support matrix reports zero
TypeError-style argument mismatches:split/array_splitforwardindices_or_sectionspositionally;transpose/swapaxes/tile/expand_dims/squeezepass axis arguments in the form each backend accepts;cumsum/nancumsum/cumprod/nancumprodemulate numpy’saxis=Noneflatten semantics via a new_maybe_flattenhelper so torch’s array-API binding stays happy;histogram/take/select/uniquerewritten to use_strip_none_kwargs+ dispatch helper;corrcoef/covdispatch unary wheny is None;nan_to_numroutes per-block via numpy on dask;eye/tripasskas a keyword and materialize a default dtype for dask auto-chunking;bincountonly forwards JAX’slength=when set.Backend-aware dtype translation (
_translate_dtype) insaiunit/_backend.pymaps numpy dtype objects to each backend’s native attribute, fixingarange,astype,tril_indices,triu_indices, and theQuantity.astypemethod on torch and ndonnx (both of which reject numpy dtypes).
Consolidated typing module#
New
saiunit/_typing.pymodule withArray,ArrayLike,ScalarOrArrayLike,DTypeLike,Shape,Axis,Axes, andPyTree. Everyjax.Array/jax.typing.DTypeLikereference in production code (annotations and docstrings) now uses these aliases.saiunit.typingre-exports the module so the public import path is unchanged.The new
ArrayLikeis narrow — it deliberately excludes bare Python scalars (bool/int/float/complex) so that.shape/.ndim/.dtypereads do not produce union-attr false positives. Approximately 658jax.typing.ArrayLikereferences acrosslax/math/fft/linalg/sparsehave been migrated.
Improvements#
Quantity core#
Quantity(mantissa, dtype=...)honorsdtypeon cupy / torch / dask / ndonnx mantissas (previously silently ignored).Quantity.mantissasetter coerces the new value to the existing backend of theQuantityrather than silently lifting torch / cupy / dask / ndonnx arrays to JAX._check_units_and_collect_values(the helper behindQuantity([q1, q2, ...])) honorsset_default_backend()/using_backend()instead of always landing on JAX when JAX is installed.Quantity.shapeprefersm.shapeforarray_api_compatbackends that do not exposexp.shape;Quantity.repeat,round,argmax,argmin,take,nanprod,expand_dims, andunsqueezeresolve their backend at call time via_xp_attrand pick the calling form each backend accepts.
Activation functions#
relu,sigmoid,softplus,gelu, … now guard withrequire_jax_backendand raise a cleanBackendErroron non-JAX inputs (instead of crashing withAttributeError: 'NoneType' object has no attribute 'relu').leaky_reluis implemented viawhereand remains backend-agnostic.Quantity.atdocstring now ships a per-backend repeated-index semantics table.
Lazy / symbolic backends#
isnan/isinf/isfinite/isreal/gradientroute throughget_backend(...)rather than hard-codedjnp.*, so non-JAX arrays no longer hit a numpy-only path.saiunit.fft.fftn/ifftn/rfftn/irfftnno longer force Python-scalar / list inputs throughjnp.asarrayfor the input-ndim probe.
Backend-specific notes#
NumPy / CuPy: repeated-index updates under
Quantity.atusenp.<op>.at/cupy.<op>.atso the JAX “all-updates-applied” semantics carry over.PyTorch:
addusesindex_put_(accumulate=True)and matches JAX for repeated indices.multiply/divide/min/max/applyuse gather + op + scatter and follow last-write-wins semantics for repeated indices (one of the rare places torch’s native semantics differ from JAX).Dask:
Quantity.atupdates are expressed viada.whereover a positional boolean mask, so the task graph stays lazy and chunked. Supported index types: scalar int, 1D integer array, slice, ellipsis, and same-shape boolean mask.ndonnx:
Quantity.at[...].set/get/...raisesBackendError;saiunit.fft.tril_indices/triu_indicesuse a static numpy fallback wrapped withxp.asarray(ndonnx exposes no nativetril_indices).
Known limitations#
modeemulation on.at[...]is best-effort on non-JAX backends for scalar-int and 1D-integer-array indices only. Slice/boolean/ellipsis indices ignoremodebecause they cannot go out of bounds against a same-shape source.daskdoes not yet support multi-dim fancy integer indexing under.at; use.to_numpy()for those patterns.After the dispatch sweep, every remaining non-pass cell in the function-level support matrix traces to a genuine missing-attribute or
NotImplementedErroron the backend rather than a saiunit-side argument mismatch. The 11 remaining ndonnx fails are real backend limitations (nocomplex128;NotImplementedErrorforcopysign/expm1/hypot/log1p/nextafter/signbit).
Documentation#
Function-level backend support matrix. New generated rst page (
docs/backends/feature_support_matrix.rst) documents per-(function, backend) support across the five locally-installed backends, covering 294 mapped + 47 non-dispatched entries insaiunit.math, all 35 insaiunit.linalg, all 16 insaiunit.fft, and 78/79Quantitymethods. Companion tooling (dev/backend_support_sweep.py,dev/backend_support_render.py,dev/backend_support_data.json) produces a byte-deterministic, reproducible sweep.Installation guide rewrite (
docs/getting_started/installation.rst). Full reference covering requirements, quick install, per-extra options, combining rules, source/editable install, install verification, upgrade/uninstall, troubleshooting (BackendError, CUDA mismatch, CuPy runtime, PyTorch wheel pitfalls), and BrainX ecosystem install. The landing page (index.rst) keeps only the two most common install commands and links out to the full reference.Per-backend Jupyter notebooks (
numpy,jax,torch,dask,ndonnx,overview) refreshed with updated outputs and markdown formatting.Chaobrain ecosystem documentation links (
brainmodeling,brainstate,brainpy-state,brainunit,braincell,brainmass,brainevent,braintrace,braintools) redirected from*.readthedocs.ioto the newbrainx.chaobrain.comdomain. Third-party RTD links remain intact.BrainUnit README header / Sphinx logo serves
brainunit.webpfrombrainx.chaobrain.com(~91% bandwidth saving). The top-level SAIUnit README continues to use its locallogo.pngidentity.
Type checking#
mypy saiunit/reports zero errors across 111 source files (down from 2,019 before the cleanup). The newtype_checkCI job blocks any regression on PRs.Targeted
# type: ignorecomments are limited to documented stub gaps (jax.laxArrayvs the narrowArrayLike,absl-pyparameterized.productlimitations in tests, missing third-party stubs forbrainstate/cupy/opt_einsum/jax.util).[[tool.mypy.overrides]]inpyproject.tomlcarriesignore_missing_imports = truefor optional backend libraries so the CItype_checkjob does not need to install ~1 GB of extras.
Packaging#
PEP 561
py.typedmarker shipped in bothsaiunitandbrainunitpackage roots and declared in package-data, so downstream type checkers (mypy, pyright) honor the inline annotations.setuptools.packages.findnow excludes thebrainunit*subtree from thesaiunitwheel — the wheel drops from ~180 to 130 entries with no brainunit references, whilesaiunit/py.typedstill ships.
CI#
Five per-backend isolation jobs (
test_pure_numpy,test_pure_jax,test_pure_torch,test_pure_dask,test_pure_ndonnx) onCI.yml. Each installs exactly one array backend on top ofrequirements.txtand runs the full suite withSAIUNIT_DEFAULT_BACKENDset, verifying thatpip install saiunit[<backend>]works end-to-end and that no other backend library leaks in.cupyis intentionally excluded — GitHub free runners have no GPU.conftest.pyreadsSAIUNIT_DEFAULT_BACKENDat session start, callsset_default_backend, and skip-collects any test file that importsjaxat module level whenHAS_JAXisFalseso the no-jax CI variants don’t crash during collection.test_pure_jaxinstalls the in-treebrainunitover the PyPI-pinned copy viamake_brainunit_setup.py+--force-reinstall --no-depsso the checkout’s API surface (after the Equinox-shim removal) is the one being tested.deploy-docs.ymldrops thepush:mainbuild-only trigger:release:releasedandworkflow_dispatchboth build and deploy the docs; the redundant build-only path on push is gone.
Removed#
The
compatible_with_equinoxtoggle,compat_with_equinoxglobal, the Equinox branch inQuantity.__pow__, the associated tests, the docs reference, and the mypy ignore overrides for the Equinox path.The
promote_integerskeyword and**kwargscatch-all fromsaiunit.math.sumandsaiunit.math.prod(see Breaking changes).
Version 0.2.2#
Highlights#
This release turns saiunit into a multi-backend library. Quantity can now
wrap NumPy, JAX, CuPy, PyTorch, Dask, and ndonnx arrays, with operations
dispatched via the array API standard (array_api_compat). JAX is now an
optional dependency — the core package installs and runs on NumPy alone.
This release also lands a 22-issue audit of Unit naming, display, hashing,
and parsing, and tightens correctness in Quantity’s hash/equality and
tracer interactions.
Breaking Changes#
JAX is no longer a mandatory dependency. Install with
pip install saiunitfor the NumPy-only build, orpip install "saiunit[jax]"(or[cpu]/[cuda12]/[cuda13]/[tpu]) to enable JAX. Without JAX, the default backend auto-selects"numpy"and JAX-only modules (saiunit.autograd,saiunit.lax,saiunit.sparse) raiseBackendErroron import with an install hint.Quantity(np.ndarray(...))now preserves the mantissa asnp.ndarrayinstead of implicitly converting tojax.Array. Call.to_jax()or usewith using_backend("jax"):to restore the previous behaviour.Quantityis now unhashable (__hash__ = None).Quantity.__eq__returns an array, so any hash implementation would violate the hash/eq invariant. This matches NumPy/JAX array semantics across all backends.Unit(dim, base=B)withB != 10now raisesValueError. Previously the non-decimal base was silently folded intofactorand then forgotten. Encode non-decimal scales directly infactor.Unit("symbol", scale=..., factor=..., name=..., ...)now raisesTypeErrorwhen extra construction kwargs are combined with a stringdim. Previously the extras were silently dropped.Unit(dim, factor=...)now raisesValueErrorfor NaN or infinite factors instead of constructing a poisoned unit that propagated NaN through all subsequent arithmetic.Quantity + UnitandUnit + Quantitynow both raiseTypeErrorsymmetrically. Previously the former silently promoted the bare Unit toQuantity(1, unit)while the latter raised — makingq + metreandmetre + qproduce different results.
New Features#
Multi-backend support#
NumPy backend.
Quantitycan wrapnp.ndarraydirectly; math, linalg, and fft operations dispatch througharray_api_compat.CuPy backend (
saiunit[cupy]). GPU arrays via the CuPy array API, withQuantity.to_cupy(device=None)for conversion.PyTorch backend (
saiunit[torch]). Torch tensors withQuantity.to_torch(device=None, dtype=None).Dask backend (
saiunit[dask]). Lazy chunked arrays viaQuantity.to_dask(chunks='auto');__repr__and other materializing methods are lazy-safe and guarded to avoid implicit computation.ndonnx backend (
saiunit[ndonnx]). Symbolic ONNX-graph arrays viaQuantity.to_ndonnx()for export-oriented workflows.saiunit[all]meta-extra installs every optional backend.
Backend control API#
Quantity.backendproperty reporting the active backend name ('numpy','jax','cupy','torch','dask','ndonnx').Quantity.to_numpy()/.to_jax()/.to_cupy()/.to_torch()/.to_dask()/.to_ndonnx()conversion methods.saiunit.set_default_backend(),saiunit.get_default_backend(), andsaiunit.using_backend()context manager for controlling the default backend when input backend is ambiguous (Python scalars, list inputs).saiunit.is_numpy_array(),is_jax_array(),is_cupy_array(),is_torch_array(),is_dask_array(), andis_ndonnx_array()detector helpers.Quantity.__array_ufunc__so calls likenp.sin(quantity)andnp.add(q1, q2)preserve units instead of stripping them.saiunit.BackendErrorexception type (subclass ofTypeError) raised by JAX-only modules (saiunit.lax,saiunit.sparse,saiunit.autograd, and the customexprelprimitive) when given a non-JAX backend, with a clear"call .to_jax() first"(or install) hint.
Unit additions#
SI-prefixed kelvin variants (
ykelvinthroughYkelvin, includingmK,uK,nK,kK,MK), bringing kelvin in line with every other base unit.parse_unit("mK")now succeeds.
Improvements#
Unit display, hashing, and parsing (22-issue audit)#
Compound-exponent normalization.
metre * metre2now displays asm^3instead ofm * m^2; display parts are merged after compound arithmetic.Eager standard-name resolution. Compound results from
__mul__/__div__/__pow__now write the resolved standard-unit name intoself._name/self._dispnameat construction time, keepingunit.name,unit.dispname, andstr(unit)in sync. Display parts survivecopy,deepcopy, and pickle round-trips.Hash/eq invariant restored.
Unit.__hash__no longer folds in spelling fields (name/dispname), so aliases likemetreandmeterhash and compare consistently — sets and dicts no longer hold silent duplicates.Unit.__eq__now compares(dim, scale, base, factor, _canonical_str()), with a newis_unit_equal_math()helper for name-agnostic math equivalence;has_same_unit()delegates to it.Built-in units win over user aliases.
add_standard_unitnow stamps a monotonic registration index; built-ins (registered during import) outrank user-added aliases, so registeringUnit(metre.dim, name="aaaa_meter")no longer hijacks the canonical metre display. Alphabetical order remains the tie-breaker for same-batch registrations.Named-dimensionless identity preserved.
radian * UNITLESS,radian ** 1, andUNITLESS * radiannow render asradinstead of collapsing to bareUnit("1");radian * radian→rad^2;radian / radian→1(genuine cancellation).Parser improvements.
parse_unitnow accepts parenthesised sub-expressions in numerators ((m * s) / A), and numeric-base tokens like"2^3"are encoded asUnit(DIMENSIONLESS, factor=8.0)rather than raising. AnonymousUnitinstances (constructed without aname) now render with parser-compatible grammar, soparse_unit(repr(u))round-trips for any anonymous unit.
Core correctness#
Tracer safety.
Quantity.update_mantissa()and__setitem__now raise a clearRuntimeErrorwhen called on a traced mantissa, instead of producing silently wrong state.JIT-safe reductions. New
_reduction_count_from_shapeand_is_concrete_zerohelpers ensure reductions stay JIT-safe, and the “0 is dimensionless” convention only applies to concrete zeros — never tracers.Foreign-tensor rejection.
require_jax_backendnow names the offending backend and rejects foreign tensors (CuPy, Torch, Dask, ndonnx) at JAX-only entry points.
Documentation#
New multi-backend user-guide sections covering NumPy, CuPy, PyTorch, Dask (lazy semantics), and ndonnx (symbolic execution).
Per-backend Jupyter notebooks demonstrating end-to-end workflows.
Updated installation docs covering the new extras layout (
jax,cpu,cuda12,cuda13,tpu,cupy,torch,dask,ndonnx,all).Sphinx
conf.pygates brainx header injection onTARGET=brainunit; docs build/deploy split (deploy on release, build-only on main push).
Dependencies#
New mandatory runtime dependencies:
array_api_compat>=1.9andopt_einsum.jaxmoved to the[jax]optional extra. Install withpip install "saiunit[jax]"(or[cpu]/[cuda12]/[cuda13]/[tpu]) to enable JAX-backed features.New optional extras:
[cupy],[torch],[dask],[ndonnx], and the[all]meta-extra.
Internal / CI#
New
_jax_compatmodule centralizes safely-degradable JAX symbols (sentinel classes, no-op decorators, NumPy fallbacks fordtypes/result_type/tree-ops) when JAX is missing.New
test_no_jaxCI job installs without JAX and verifies imports plusBackendErrorgates;_no_jax_test.pysmoke suite added.CI now installs CPU-only PyTorch wheels (
--index-url https://download.pytorch.org/whl/cpu) to avoid multi-GB CUDA pulls on GPU-less runners.CI matrix extended to install
cupy/torch/dask/ndonnxon all platforms so the new backends are exercised in regression tests.brainunitlegacy package now re-exports the new backend API (BackendError,is_unit_equal_math,set_default_backend,using_backend, backend detectors).Dependency bumps:
actions/checkout4→6,actions/setup-python5→6,actions/download-artifact5→8,appleboy/ssh-action1.2.0→1.2.5,appleboy/scp-action0.1.7→1.0.0,sphinx≥8.1.3,sphinx-book-theme≥1.1,sphinx-copybutton≥0.5.2,jupyter-sphinx≥0.5.3.
Version 0.2.1#
Breaking Changes#
Removed
BlockCSRandBlockELLsparse classes: These experimental block-sparse matrix implementations (along with their benchmarks and tests) have been removed. Users should useCOO,CSR, orCSCinstead.SparseMatrixno longer inherits fromJAXSparse: The base class is now standalone, with its ownshape,size,ndim,T,block_until_ready,tree_flatten,tree_unflatten,transpose, andtodenseinterface.
Improvements#
Forward-compatible
**kwargsacross all wrapped functions: All unit-aware wrapper functions inmath,lax,linalg, andfftmodules now accept and forward**kwargsto the underlying JAX functions. This ensures compatibility with new keyword arguments added in future JAX releases without requiring saiunit updates.saiunit.math:concatenate,stack,vstack,hstack,dstack,column_stack,block,append,split,array_split,dsplit,hsplit,vsplit,tile,repeat,sort,argsort,unique,searchsorted,where,clip,interp, and many moresaiunit.lax:cond,switch,scan,while_loop,fori_loop,sort,top_k,broadcasted_iota,concatenate,conv,pad,slice,dynamic_slice,gather,scatter, and many moresaiunit.linalg:svd,cholesky,eig,eigh,eigvalsh,qr,lu,solve,det,norm,matrix_power,cross,tensordot, etc.saiunit.fft:fft,ifft,fft2,ifft2,fftn,ifftn,rfft,irfft,rfft2,irfft2,rfftn,irfftn,fftshift,ifftshift
Standalone
SparseMatrixbase class: Decoupled fromjax.experimental.sparse.JAXSparseto reduce external coupling and provide a self-contained sparse matrix interface with properties (size,ndim,nse,dtype) and methods (__repr__,__len__,block_until_ready).Improved validation in sparse classes: Replaced
assertstatements with descriptiveValueErrorexceptions inCOO.with_data,CSR.with_data, andCSC.with_datafor shape, dtype, and unit mismatches.Broader sparse type checking:
isinstancechecks in binary and matmul operations now accept bothJAXSparseandSparseMatrix, ensuring correct behavior after the inheritance change.Fixed
CSC.tree_unflattenerror message: Corrected the error message from"CSR.tree_unflatten"to"CSC.tree_unflatten".Explicit attribute assignment in
tree_unflatten:CSRandCSCnow setshape,indices, andindptrexplicitly instead of using__dict__.update, improving clarity and avoiding potential issues.
Version 0.2.0#
Highlights#
This release introduces unit-aware type annotations, string-based unit parsing, enhanced Matplotlib integration, and a comprehensive overhaul of error handling and unit display semantics.
New Features#
Unit-aware type annotations (
saiunit.typing): AddedQuantityLike,UnitLike, and related type aliases usingtyping.Annotated(PEP 593) for expressing physical-unit constraints in Python type hints.String-based unit parsing:
Quantitynow accepts string unit specifications during initialization (e.g.,Quantity(1.0, "meter")).Matplotlib
QuantityConverter: Full integration with Matplotlib’s unit conversion framework, enabling direct plotting ofQuantityobjects with automatic axis labeling and unit display.unit_to_scaleparameter for activation functions: Activation functions now accept an optionalunit_to_scaleparameter for explicit unit conversion.symmetrize_inputparameter: Added tocholesky,eigvalsh, andsvdfor optional input symmetrization before decomposition.amualias: Addedamuas an alias foratomic_mass/u/um_u.concrete_or_errorshim: Compatibility shim forjax.core.concrete_or_errorto maintain support across JAX versions.FFT
shapeparameter:_calculate_fftn_dimensionnow supports ashapeparameter for explicit output shape specification.
Improvements#
Unified unit display format: Refactored
display_in_unitand unit representation methods for consistent, human-readable output. Normalized exponent representation and improved formatting for dimensionless units. Removed thepython_codeparameter in favor of unified display.Error handling overhaul: Replaced
assertstatements with properTypeErrorandValueErrorexceptions across the codebase (Quantity, einops, lax, FFT, and unit-related functions) for clearer, more informative error messages.Updated physical constants (CODATA 2018):
atomic_mass,electron_volt,light_year,atmosphere,acre,fluid_ounce_imp,Btu_th,speed_of_sound, andIMFnow use more accurate conversion factors.Improved display names:
survey_foot/survey_milenow show"US survey ft"/"US survey mi";gallon_impshows"imp gal";fluid_ounce_impshows"imp fl oz";monthshows"mon";Btu_ITshows"Btu".Unit preference scoring: Standard unit retrieval now uses preference scoring for aliases, ensuring compound unit representations maintain grouping.
Removed
iscompoundattribute from theUnitclass, simplifying the internal representation.Cumulative product functions: Enhanced handling for unit-aware quantities.
Jacobian and vector gradient: Refactored to use
_argnums_partialfor improved multi-argument handling; added tests for list-styleargnums.Sparse matrix improvements: Refactored unit handling in
BlockCSRandBlockELL; addedtransposemethod.Module organization: Restructured imports across
__init__.pyand submodules for consistency; defined__all__for decorators and constants.Type annotations modernized: Updated
Union[A, B]toA | Bsyntax throughout the codebase.Documentation: Updated docstrings with examples for unit-aware functions; added installation instructions for CUDA and TPU; refreshed all Jupyter notebook examples.
Copyright updated to BrainX Ecosystem Limited.
Bug Fixes#
Fixed issue #17 (unit display edge case).
Fixed cumulative product operations for unit-aware quantities.
Fixed error handling in tests: corrected expected exception types from
AssertionErrortoTypeErrorfor invalid input cases.
Internal / CI#
Refactored version handling and updated main entry point structure.
Removed deprecated JAX version testing from CI configuration.
Added
brainstateto optional testing dependencies.Updated CI configurations to include BrainUnit installation steps.
Removed
sys.version_infochecks for Python version compatibility.
Version 0.1.4#
Added numerically stable
exprelfunction with comprehensive test coverageUpdated lax array creation to use
jax.numpyfor zero initializationUpdated CI JAX version for improved compatibility
Improved code quality and removed redundant tests
Version 0.1.3#
Compatible with
jax>=0.8.2
Version 0.1.2#
Renamed
CustomArray.valuetoCustomArray.datafor API consistencyStreamlined math unwrapping for improved performance
Refactored math module for better Quantity/CustomArray support
Added dtype aliases for convenience
Fixed matplotlib convert for zero-sized inputs
Registered Array class as a PyTree node for JAX compatibility
Added support for Python 3.14
Updated CI configuration and dependencies
Version 0.1.1#
Fixed dimension and unit checks in convert method to handle empty input
Bug fixes and stability improvements
Version 0.1.0#
Introduced
CustomArrayclass and integrated it across saiunit modulesAdded
Arrayclass inheriting fromCustomArrayAdded
maybe_custom_arrayandmaybe_custom_array_treeutilities for type checkingAdded comprehensive unit tests for CustomArray integration
Improved support for einops, activation functions, FFT, and linear algebra operations
Enhanced Celsius conversion functions with CustomArray compatibility
Added tutorial and documentation for CustomArray
Version 0.0.19#
Added
CustomArrayclass as the foundation for custom array typesRefactored activation functions to support CustomArray
Added
gatherfunctionRefined
math,linalg,autograd, and constants modulesEnabled Quantity hashing with
__hash__methodFixed interp function and adjusted unit handling in output
Moved metadata from setup.py to pyproject.toml
Version 0.0.2#
The new version of saiunit, which separates the Quantity into the mantissa and the unit.
This design is more flexible and allows for more complex operations, enabling to represent the very large or very small values.
Version 0.0.1#
The first release of the project.