ndonnx backend

ndonnx backend#

ndonnx is an array-API-compatible library that builds an ONNX graph symbolically instead of computing eagerly. Use it when you want to build an ONNX model whose computation involves unit-aware quantities.

Installation#

pip install brainunit[ndonnx]

Graceful import#

import brainunit as u

try:
    import ndonnx
    HAVE_NDONNX = True
    print('ndonnx version:', ndonnx.__version__)
except ImportError:
    HAVE_NDONNX = False
    print('ndonnx not installed; install with: pip install brainunit[ndonnx]')

ndonnx version: 0.21.0

Quick start — symbolic by default#

if HAVE_NDONNX:
    import numpy as np
    q = u.Quantity(ndonnx.asarray(np.array([1.0, 2.0, 3.0])), unit=u.meter)
    print('backend       =', q.backend)
    print('(q + q).backend =', (q + q).backend)

backend       = ndonnx
(q + q).backend = ndonnx

Unit checking is independent of the backend#

Dimensional analysis lives on the Python Quantity. Mixing units fails the same way regardless of mantissa kind.

if HAVE_NDONNX:
    import numpy as np
    from brainunit import UnitMismatchError
    a = u.Quantity(ndonnx.asarray(np.array([1.0])), unit=u.meter)
    b = u.Quantity(ndonnx.asarray(np.array([1.0])), unit=u.second)
    try:
        a + b
    except UnitMismatchError as exc:
        print('expected:', exc)

expected: Cannot calculate 
[array(data: 1.0, dtype=float64)] m + [array(data: 1.0, dtype=float64)] s, because units do not match: m != s

Materialization#

Quantity.to_numpy() calls ndonnx’s unwrap_numpy on the mantissa rather than np.asarray (which would give you back a 0-d object wrapper).

if HAVE_NDONNX:
    import numpy as np
    q = u.Quantity(ndonnx.asarray(np.array([1.0, 2.0])), unit=u.meter)
    eager = q.to_numpy()
    print(eager.backend, eager.mantissa)

numpy [1. 2.]

Caveats#

ndonnx is still maturing; some brainunit.math / brainunit.linalg operations may not have an ndonnx implementation. When that happens, the ndonnx error surfaces unwrapped — brainunit does not catch it. Consult the ndonnx documentation for the supported op set.
brainunit does not provide brainunit-level ONNX export helpers. Unit information lives on the Python Quantity object and is not encoded in the ONNX graph. Export the mantissa using ndonnx’s own workflow.

ndonnx backend

Contents