saiunit 0.2.0__tar.gz → 0.2.2__tar.gz

{saiunit-0.2.0 → saiunit-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: saiunit
-Version: 0.2.0
+Version: 0.2.2
 Summary: Enabling Unit-aware Computations for AI-driven Scientific Computing.
 Author-email: SAIUnit Developers <chao.brain@qq.com>
 License-Expression: Apache-2.0
@@ -28,12 +28,15 @@ Classifier: Topic :: Software Development :: Libraries
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: jax
 Requires-Dist: numpy
 Requires-Dist: typing_extensions
+Requires-Dist: array_api_compat>=1.9
+Requires-Dist: opt_einsum
 Provides-Extra: testing
 Requires-Dist: pytest; extra == "testing"
 Requires-Dist: brainstate; extra == "testing"
+Provides-Extra: jax
+Requires-Dist: jax; extra == "jax"
 Provides-Extra: cpu
 Requires-Dist: jax[cpu]; extra == "cpu"
 Provides-Extra: cuda12
@@ -42,6 +45,16 @@ Provides-Extra: cuda13
 Requires-Dist: jax[cuda13]; extra == "cuda13"
 Provides-Extra: tpu
 Requires-Dist: jax[tpu]; extra == "tpu"
+Provides-Extra: cupy
+Requires-Dist: cupy-cuda12x>=13.0; extra == "cupy"
+Provides-Extra: torch
+Requires-Dist: torch>=2.0; extra == "torch"
+Provides-Extra: dask
+Requires-Dist: dask[array]>=2024.1; extra == "dask"
+Provides-Extra: ndonnx
+Requires-Dist: ndonnx>=0.9; extra == "ndonnx"
+Provides-Extra: all
+Requires-Dist: saiunit[cupy,dask,jax,ndonnx,torch]; extra == "all"
 Dynamic: license-file
 
 <p align="center">
@@ -124,17 +137,71 @@ jax.vmap(f)(u.math.arange(0. * u.mV, 10. * u.mV, 1. * u.mV))
 
 
 
+## Multiple-backend support
+
+`saiunit` is **backend-agnostic**: a `Quantity` pairs a unit with an array
+mantissa, and that mantissa can live on any of the supported array libraries.
+Every unit-aware operation dispatches to the matching backend, so you can stay
+in one library end-to-end or convert with a single method call.
+
+| Backend  | Mantissa            | Install                  | When to use                                 |
+|----------|---------------------|--------------------------|---------------------------------------------|
+| `numpy`  | `numpy.ndarray`     | core (always installed)  | eager CPU, scipy/pandas/sklearn interop     |
+| `jax`    | `jax.Array`         | `saiunit[jax]` (or `[cpu]`/`[cuda12]`/`[cuda13]`/`[tpu]`) | autograd, JIT, vmap, accelerators           |
+| `cupy`   | `cupy.ndarray`      | `saiunit[cupy]`          | NVIDIA GPU arrays                           |
+| `torch`  | `torch.Tensor`      | `saiunit[torch]`         | PyTorch models, torch autograd              |
+| `dask`   | `dask.array.Array`  | `saiunit[dask]`          | out-of-core / parallel, lazy compute        |
+| `ndonnx` | `ndonnx.Array`      | `saiunit[ndonnx]`        | symbolic graph for ONNX export              |
+
+Select or override the backend explicitly with `u.using_backend(...)` or
+`u.set_default_backend(...)`; convert with `q.to_jax()` / `q.to_numpy()` /
+`q.to_cupy()` / `q.to_torch()` / `q.to_dask()` / `q.to_ndonnx()`. Requesting an
+uninstalled backend raises `saiunit.BackendError` with the install command —
+not a bare `ImportError`. See the
+[Backends documentation](https://saiunit.readthedocs.io/en/latest/backends/overview.html)
+for the full story.
+
+
 ## Installation
 
-``saiunit`` has been well tested on ``python>=3.9`` + ``jax>=0.4.30`` environments, and can be installed on Windows, Linux, and MacOS.
+``saiunit`` has been well tested on ``python>=3.10`` and can be installed on
+Windows, Linux, and MacOS. The core package depends only on NumPy. JAX is
+optional — install it to enable the ``saiunit.autograd``, ``saiunit.lax``,
+and ``saiunit.sparse`` submodules, the custom ``exprel`` primitive, and the
+``"jax"`` backend.
 
-You can install ``saiunit`` via pip:
+Install the NumPy-only core:
 
 ```bash
 pip install saiunit --upgrade
 ```
 
-which should install in about 1 minute. If you want to install the latest version from the source, you can clone the repository and install it:
+Or pull in JAX with the accelerator build that matches your hardware:
+
+```bash
+pip install -U saiunit[jax]      # plain JAX
+pip install -U saiunit[cpu]      # pinned JAX CPU wheels
+pip install -U saiunit[cuda12]   # JAX on CUDA 12
+pip install -U saiunit[cuda13]   # JAX on CUDA 13
+pip install -U saiunit[tpu]      # JAX on TPU
+```
+
+Opt into additional array backends with the matching extra:
+
+```bash
+pip install -U saiunit[cupy]     # CuPy (NVIDIA GPU)
+pip install -U saiunit[torch]    # PyTorch
+pip install -U saiunit[dask]     # Dask
+pip install -U saiunit[ndonnx]   # ndonnx
+pip install -U saiunit[all]      # jax + cupy + torch + dask + ndonnx
+```
+
+Without JAX, the NumPy backend is auto-selected and any access to a
+JAX-only submodule (``saiunit.autograd``, ``saiunit.lax``, ``saiunit.sparse``)
+raises ``saiunit.BackendError`` with an install hint. The optional extras are
+independent and can be combined freely.
+
+To install the latest version from source:
 
 ```bash
 git clone https://github.com/chaobrain/saiunit.git

{saiunit-0.2.0 → saiunit-0.2.2}/README.md

@@ -78,17 +78,71 @@ jax.vmap(f)(u.math.arange(0. * u.mV, 10. * u.mV, 1. * u.mV))
 
 
 
+## Multiple-backend support
+
+`saiunit` is **backend-agnostic**: a `Quantity` pairs a unit with an array
+mantissa, and that mantissa can live on any of the supported array libraries.
+Every unit-aware operation dispatches to the matching backend, so you can stay
+in one library end-to-end or convert with a single method call.
+
+| Backend  | Mantissa            | Install                  | When to use                                 |
+|----------|---------------------|--------------------------|---------------------------------------------|
+| `numpy`  | `numpy.ndarray`     | core (always installed)  | eager CPU, scipy/pandas/sklearn interop     |
+| `jax`    | `jax.Array`         | `saiunit[jax]` (or `[cpu]`/`[cuda12]`/`[cuda13]`/`[tpu]`) | autograd, JIT, vmap, accelerators           |
+| `cupy`   | `cupy.ndarray`      | `saiunit[cupy]`          | NVIDIA GPU arrays                           |
+| `torch`  | `torch.Tensor`      | `saiunit[torch]`         | PyTorch models, torch autograd              |
+| `dask`   | `dask.array.Array`  | `saiunit[dask]`          | out-of-core / parallel, lazy compute        |
+| `ndonnx` | `ndonnx.Array`      | `saiunit[ndonnx]`        | symbolic graph for ONNX export              |
+
+Select or override the backend explicitly with `u.using_backend(...)` or
+`u.set_default_backend(...)`; convert with `q.to_jax()` / `q.to_numpy()` /
+`q.to_cupy()` / `q.to_torch()` / `q.to_dask()` / `q.to_ndonnx()`. Requesting an
+uninstalled backend raises `saiunit.BackendError` with the install command —
+not a bare `ImportError`. See the
+[Backends documentation](https://saiunit.readthedocs.io/en/latest/backends/overview.html)
+for the full story.
+
+
 ## Installation
 
-``saiunit`` has been well tested on ``python>=3.9`` + ``jax>=0.4.30`` environments, and can be installed on Windows, Linux, and MacOS.
+``saiunit`` has been well tested on ``python>=3.10`` and can be installed on
+Windows, Linux, and MacOS. The core package depends only on NumPy. JAX is
+optional — install it to enable the ``saiunit.autograd``, ``saiunit.lax``,
+and ``saiunit.sparse`` submodules, the custom ``exprel`` primitive, and the
+``"jax"`` backend.
 
-You can install ``saiunit`` via pip:
+Install the NumPy-only core:
 
 ```bash
 pip install saiunit --upgrade
 ```
 
-which should install in about 1 minute. If you want to install the latest version from the source, you can clone the repository and install it:
+Or pull in JAX with the accelerator build that matches your hardware:
+
+```bash
+pip install -U saiunit[jax]      # plain JAX
+pip install -U saiunit[cpu]      # pinned JAX CPU wheels
+pip install -U saiunit[cuda12]   # JAX on CUDA 12
+pip install -U saiunit[cuda13]   # JAX on CUDA 13
+pip install -U saiunit[tpu]      # JAX on TPU
+```
+
+Opt into additional array backends with the matching extra:
+
+```bash
+pip install -U saiunit[cupy]     # CuPy (NVIDIA GPU)
+pip install -U saiunit[torch]    # PyTorch
+pip install -U saiunit[dask]     # Dask
+pip install -U saiunit[ndonnx]   # ndonnx
+pip install -U saiunit[all]      # jax + cupy + torch + dask + ndonnx
+```
+
+Without JAX, the NumPy backend is auto-selected and any access to a
+JAX-only submodule (``saiunit.autograd``, ``saiunit.lax``, ``saiunit.sparse``)
+raises ``saiunit.BackendError`` with an install hint. The optional extras are
+independent and can be combined freely.
+
+To install the latest version from source:
 
 ```bash
 git clone https://github.com/chaobrain/saiunit.git

{saiunit-0.2.0 → saiunit-0.2.2}/brainunit/brainunit/__init__.py

@@ -18,6 +18,7 @@ __version__ = saiunit.__version__
 __version_info__ = saiunit.__version_info__
 
 from . import autograd
+from saiunit._matplotlib_compat import enable_matplotlib_support
 from . import constants
 from . import fft
 from . import lax
@@ -34,6 +35,18 @@ from ._base_dimension import (
     get_dim_for_display,
     get_or_create_dimension,
 )
+from ._backend import (
+    get_default_backend,
+    is_cupy_array,
+    is_dask_array,
+    is_jax_array,
+    is_ndonnx_array,
+    is_numpy_array,
+    is_torch_array,
+    set_default_backend,
+    using_backend,
+)
+from ._exceptions import BackendError
 from ._base_getters import (
     array_with_unit,
     assert_quantity,
@@ -48,6 +61,7 @@ from ._base_getters import (
     have_same_dim,
     is_dimensionless,
     is_scalar_type,
+    is_unit_equal_math,
     is_unitless,
     maybe_decimal,
     split_mantissa_unit,
@@ -55,7 +69,7 @@ from ._base_getters import (
 )
 from ._base_quantity import Quantity, compatible_with_equinox
 from ._base_unit import UNITLESS, Unit, add_standard_unit, parse_unit
-from ._celsius import celsius2kelvin, kelvin2celsius
+from ._celsius import celsius2kelvin, kelvin2celsius, fahrenheit2kelvin, kelvin2fahrenheit
 from ._misc import maybe_custom_array, maybe_custom_array_tree
 from ._unit_common import *
 from ._unit_common import __all__ as _common_all
@@ -100,6 +114,18 @@ __all__ = [
               'DIMENSIONLESS',
               'DimensionMismatchError',
               'UnitMismatchError',
+              'BackendError',
+
+              # _backend
+              'get_default_backend',
+              'set_default_backend',
+              'using_backend',
+              'is_jax_array',
+              'is_numpy_array',
+              'is_cupy_array',
+              'is_torch_array',
+              'is_dask_array',
+              'is_ndonnx_array',
               'get_or_create_dimension',
               'get_dim_for_display',
 
@@ -125,6 +151,7 @@ __all__ = [
               'assert_quantity',
               'have_same_dim',
               'has_same_unit',
+              'is_unit_equal_math',
               'unit_scale_align_to_first',
               'array_with_unit',
 
@@ -140,6 +167,11 @@ __all__ = [
               # _celsius
               'celsius2kelvin',
               'kelvin2celsius',
+              'fahrenheit2kelvin',
+              'kelvin2fahrenheit',
+
+              # _matplotlib_compat
+              'enable_matplotlib_support',
 
               # old version compatibility
               'avogadro_constant',

saiunit-0.2.2/brainunit/brainunit/_backend.py

@@ -0,0 +1,17 @@
+# Copyright 2026 BrainX Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from saiunit._backend import *
+from saiunit._backend import __all__

saiunit-0.2.2/brainunit/brainunit/_exceptions.py

@@ -0,0 +1,17 @@
+# Copyright 2026 BrainX Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from saiunit._exceptions import *
+from saiunit._exceptions import __all__

{saiunit-0.2.0 → saiunit-0.2.2}/brainunit/brainunit/linalg/__init__.py

@@ -17,9 +17,13 @@ from ._linalg_change_unit import *
 from ._linalg_change_unit import __all__ as _linalg_change_unit_all
 from ._linalg_keep_unit import *
 from ._linalg_keep_unit import __all__ as _linalg_keep_unit_all
+from ._linalg_remove_unit import *
+from ._linalg_remove_unit import __all__ as _linalg_remove_unit_all
 
 __all__ = (_linalg_change_unit_all +
-              _linalg_keep_unit_all)
+              _linalg_keep_unit_all +
+              _linalg_remove_unit_all)
 
 del (_linalg_change_unit_all,
-     _linalg_keep_unit_all)
+     _linalg_keep_unit_all,
+     _linalg_remove_unit_all)

{saiunit-0.2.0 → saiunit-0.2.2}/pyproject.toml

@@ -47,25 +47,35 @@ classifiers = [
 keywords = ['physical unit', 'physical quantity', 'scientific computing', 'AI for science', ]
 
 dependencies = [
-    'jax',
     'numpy',
     'typing_extensions',
+    'array_api_compat>=1.9',
+    'opt_einsum',
 ]
 
 dynamic = ['version']
 
+
 [tool.flit.module]
 name = "saiunit"
 
+
 [project.urls]
 homepage = 'https://github.com/chaobrain/saiunit'
 repository = 'https://github.com/chaobrain/saiunit'
 "Bug Tracker" = "https://github.com/chaobrain/saiunit/issues"
 "Documentation" = "https://saiunit.readthedocs.io/"
 
+
 [project.optional-dependencies]
 testing = ['pytest', 'brainstate']
+jax = ["jax"]
 cpu = ["jax[cpu]"]
 cuda12 = ["jax[cuda12]"]
 cuda13 = ["jax[cuda13]"]
 tpu = ["jax[tpu]"]
+cupy = ["cupy-cuda12x>=13.0"]
+torch = ["torch>=2.0"]
+dask = ["dask[array]>=2024.1"]
+ndonnx = ["ndonnx>=0.9"]
+all = ["saiunit[jax,cupy,torch,dask,ndonnx]"]

{saiunit-0.2.0 → saiunit-0.2.2}/saiunit/__init__.py

@@ -14,13 +14,15 @@
 # ==============================================================================
 
 """
-saiunit -- Physical units for JAX arrays.
+saiunit -- Physical units for JAX and NumPy arrays.
 
-``saiunit`` provides a :class:`Quantity` type that pairs a JAX array with a
-physical :class:`Unit`, ensuring dimensional correctness at every arithmetic
-operation.  It also supplies the standard SI base and derived units (e.g.
-``meter``, ``second``, ``volt``), physical constants, and unit-aware wrappers
-for NumPy/JAX math functions.
+``saiunit`` provides a :class:`Quantity` type that pairs a JAX array or
+NumPy array with a physical :class:`Unit`, ensuring dimensional correctness
+at every arithmetic operation. The backend is detected from the mantissa
+type; users can force a default with :func:`set_default_backend` or the
+:func:`using_backend` context manager. ``saiunit`` also supplies the standard
+SI base and derived units (e.g. ``meter``, ``second``, ``volt``), physical
+constants, and unit-aware wrappers for NumPy/JAX math functions.
 
 Subpackages
 -----------
@@ -52,15 +54,13 @@ Examples
     True
 """
 
-from . import _matplotlib_compat
-from . import autograd
+from ._matplotlib_compat import enable_matplotlib_support
 from . import constants
 from . import fft
-from . import lax
 from . import linalg
 from . import math
-from . import sparse
 from . import typing
+from ._jax_compat import HAS_JAX as _HAS_JAX
 from ._base_decorators import assign_units, check_dims, check_units
 from ._base_dimension import (
     DIMENSIONLESS,
@@ -70,6 +70,18 @@ from ._base_dimension import (
     get_dim_for_display,
     get_or_create_dimension,
 )
+from ._backend import (
+    get_default_backend,
+    is_cupy_array,
+    is_dask_array,
+    is_jax_array,
+    is_ndonnx_array,
+    is_numpy_array,
+    is_torch_array,
+    set_default_backend,
+    using_backend,
+)
+from ._exceptions import BackendError
 from ._base_getters import (
     array_with_unit,
     assert_quantity,
@@ -84,6 +96,7 @@ from ._base_getters import (
     have_same_dim,
     is_dimensionless,
     is_scalar_type,
+    is_unit_equal_math,
     is_unitless,
     maybe_decimal,
     split_mantissa_unit,
@@ -91,7 +104,7 @@ from ._base_getters import (
 )
 from ._base_quantity import Quantity, compatible_with_equinox
 from ._base_unit import UNITLESS, Unit, add_standard_unit, parse_unit
-from ._celsius import celsius2kelvin, kelvin2celsius
+from ._celsius import celsius2kelvin, kelvin2celsius, fahrenheit2kelvin, kelvin2fahrenheit
 from ._misc import maybe_custom_array, maybe_custom_array_tree
 from ._unit_common import *
 from ._unit_common import __all__ as _common_all
@@ -137,6 +150,18 @@ __all__ = [
               'DIMENSIONLESS',
               'DimensionMismatchError',
               'UnitMismatchError',
+              'BackendError',
+
+              # _backend
+              'get_default_backend',
+              'set_default_backend',
+              'using_backend',
+              'is_jax_array',
+              'is_numpy_array',
+              'is_cupy_array',
+              'is_torch_array',
+              'is_dask_array',
+              'is_ndonnx_array',
               'get_or_create_dimension',
               'get_dim_for_display',
 
@@ -162,6 +187,7 @@ __all__ = [
               'assert_quantity',
               'have_same_dim',
               'has_same_unit',
+              'is_unit_equal_math',
               'unit_scale_align_to_first',
               'array_with_unit',
 
@@ -177,6 +203,11 @@ __all__ = [
               # _celsius
               'celsius2kelvin',
               'kelvin2celsius',
+              'fahrenheit2kelvin',
+              'kelvin2fahrenheit',
+
+              # _matplotlib_compat
+              'enable_matplotlib_support',
 
               # old version compatibility
               'avogadro_constant',
@@ -189,4 +220,36 @@ __all__ = [
               'magnetic_constant',
               'molar_mass_constant',
           ] + _common_all + _std_units_all + _constants_all
-del _common_all, _std_units_all, _matplotlib_compat, _constants_all
+del _common_all, _std_units_all, _constants_all
+
+# ---------------------------------------------------------------------------
+# Lazy submodule loading for JAX-only features.
+#
+# ``autograd``, ``lax`` and ``sparse`` use JAX primitives that have no NumPy
+# equivalent. Importing them eagerly would force a hard JAX dependency on
+# every ``import saiunit`` call. Instead, expose them through a module-level
+# ``__getattr__``: the first attribute access triggers the real import and,
+# if JAX is missing, raises :class:`~saiunit._exceptions.BackendError` with
+# an actionable install hint.
+# ---------------------------------------------------------------------------
+
+_JAX_ONLY_SUBMODULES = ("autograd", "lax", "sparse")
+__all__ = __all__ + ["autograd", "lax", "sparse"]
+
+
+def __getattr__(name):
+    if name in _JAX_ONLY_SUBMODULES:
+        if not _HAS_JAX:
+            from ._exceptions import BackendError
+            raise BackendError(
+                f"saiunit.{name} requires JAX. Install with: pip install saiunit[jax]"
+            )
+        import importlib
+        mod = importlib.import_module(f"saiunit.{name}")
+        globals()[name] = mod
+        return mod
+    raise AttributeError(f"module 'saiunit' has no attribute {name!r}")
+
+
+def __dir__():
+    return sorted(set(__all__) | set(globals()))