ndonnx

An ONNX-backed array library that is compliant with the Array API standard.

Installation

Releases are available on PyPI and conda-forge.

# using pip
pip install ndonnx
# using conda
conda install ndonnx
# using pixi
pixi add ndonnx

Development

You can install the package in development mode using:

git clone https://github.com/quantco/ndonnx
cd ndonnx

# For Array API tests
git submodule update --init --recursive

pixi shell
pre-commit run -a
pip install --no-build-isolation --no-deps -e .
pytest tests -n auto

Quick start

ndonnx is an ONNX based python array library.

It has a couple of key features:

It implements the Array API standard. Standard compliant code can be executed without changes across numerous backends such as like NumPy, JAX and now ndonnx.

import numpy as np
import ndonnx as ndx
import jax.numpy as jnp

def mean_drop_outliers(a, low=-5, high=5):
    xp = a.__array_namespace__()
    return xp.mean(a[(low < a) & (a < high)])

np_result = mean_drop_outliers(np.asarray([-10, 0.5, 1, 5]))
jax_result = mean_drop_outliers(jnp.asarray([-10, 0.5, 1, 5]))
onnx_result = mean_drop_outliers(ndx.asarray([-10, 0.5, 1, 5]))

assert np_result == onnx_result.to_numpy() == jax_result == 0.75

It supports ONNX export. This allows you persist your logic into an ONNX computation graph.

import ndonnx as ndx
import onnx

# Instantiate placeholder ndonnx array
x = ndx.array(shape=("N",), dtype=ndx.float32)
y = mean_drop_outliers(x)

# Build and save ONNX model to disk
model = ndx.build({"x": x}, {"y": y})
onnx.save(model, "mean_drop_outliers.onnx")

You can then make predictions using a runtime of your choice.

import onnxruntime as ort
import numpy as np

inference_session = ort.InferenceSession("mean_drop_outliers.onnx")
prediction, = inference_session.run(None, {
    "x": np.array([-10, 0.5, 1, 5], dtype=np.float32),
})
assert prediction == 0.75

In the future we will be enabling a stable API for an extensible data type system. This will allow users to define their own data types and operations on arrays with these data types.

Array API coverage

Array API compatibility tested against the official array-api-tests suite.
Missing coverage is tracked in the skips.txt file.
Contributions are welcome!

Run the tests with:

pixi run arrayapitests

Dependencies of ndonnx (3).
Dependency	Optional
numpy	false
spox	false
typing_extensions	false

Package profile

ndonnx

Releases

Release table

PyPI Downloads

Dependencies

Details

Project URLs

Dependencies

Package description