Defining ONNX graphs with ndonnx

submited by

Style Pass

2024-10-05 13:30:05

ndonnx is an ONNX-backed Python array library that implements the Array API standard. It helps us take our machine learning models into production at QuantCo by facilitating ONNX export of Array API compliant code and providing a high level API for constructing ONNX graphs.

In the previous post in our ONNX blog series, we showed how you can export a trained linear regression model to ONNX using the Spox library. You might have noticed that we implemented the prediction path twice - once using NumPy (LinearRegression.predict) for training and experimentation and a second time using Spox (linear_regression) in order to be able to export model to ONNX.

Implementing the same algorithm twice using two completely different primitives is error-prone and challenging to maintain. It forces library developers to learn a second toolchain just to support ONNX export and continuously maintain two code paths that need to remain semantically equivalent, detracting from feature development time. Unfortunately, manually reimplementing code in this style is still the state-of-the-art approach in the broader ONNX community, with libraries like scikit-learn having dedicated libraries for ONNX conversion1.

We identified that the root cause of this problem is the lack of interoperability between ONNX tooling and existing array libraries.