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
except ImportError:
    HAVE_NDONNX = False
    print('ndonnx not installed; install with: pip install brainunit[ndonnx]')

ndonnx not installed; install with: pip install brainunit[ndonnx]

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)

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)

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)

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