NumPy function#

This page answers common “how do I…” questions for exporting a Python function that processes numpy arrays into ONNX.

How to export a function processing numpy arrays#

Use yobx.sql.trace_numpy_to_onnx() with a representative sample input. Only the input dtype and shape are used during export.

<<<

import numpy as np
from yobx.helpers.onnx_helper import pretty_onnx
from yobx.sql import trace_numpy_to_onnx


def normalize_rows(X):
    norms = np.sqrt(np.sum(X ** np.float32(2), axis=1, keepdims=True))
    safe_norms = np.where(norms < np.float32(1e-8), np.float32(1), norms)
    return X / safe_norms


X_sample = np.zeros((3, 4), dtype=np.float32)
onx = trace_numpy_to_onnx(normalize_rows, X_sample)
print(pretty_onnx(onx))

>>>

    opset: domain='' version=21
    opset: domain='ai.onnx.ml' version=1
    input: name='X' type=dtype('float32') shape=['batch', 4]
    init: name='init1_s_' type=float32 shape=() -- array([2.], dtype=float32)-- Opset.make_node.1/Small
    init: name='init7_s1_1' type=int64 shape=(1,) -- array([1])           -- Opset.make_node.1/Shape
    init: name='init1_s_2' type=float32 shape=() -- array([1.e-08], dtype=float32)-- Opset.make_node.1/Small
    init: name='init1_s_3' type=float32 shape=() -- array([1.], dtype=float32)-- Opset.make_node.1/Small
    Pow(X, init1_s_) -> _onx_pow_X
      ReduceSum(_onx_pow_X, init7_s1_1, keepdims=1) -> _onx_reducesum_pow_X
        Sqrt(_onx_reducesum_pow_X) -> _onx_sqrt_reducesum_pow_X
          Less(_onx_sqrt_reducesum_pow_X, init1_s_2) -> _onx_less_sqrt_reducesum_pow_X
          Where(_onx_less_sqrt_reducesum_pow_X, init1_s_3, _onx_sqrt_reducesum_pow_X) -> _onx_where_less_sqrt_reducesum_pow_X
            Div(X, _onx_where_less_sqrt_reducesum_pow_X) -> _onx_div_X
    output: name='_onx_div_X' type='NOTENSOR' shape=None

How to export a function with multiple inputs and outputs#

Use input_names and output_names when your function takes more than one input array and returns multiple arrays.

<<<

import numpy as np
from yobx.helpers.onnx_helper import pretty_onnx
from yobx.sql import trace_numpy_to_onnx


def combine(A, B):
    return A + B, A - B


A = np.random.randn(2, 3).astype(np.float32)
B = np.random.randn(2, 3).astype(np.float32)
onx = trace_numpy_to_onnx(
    combine,
    A,
    B,
    input_names=["A", "B"],
    output_names=["sum", "diff"],
)
print(pretty_onnx(onx))

>>>

    opset: domain='' version=21
    opset: domain='ai.onnx.ml' version=1
    input: name='A' type=dtype('float32') shape=['batch', 3]
    input: name='B' type=dtype('float32') shape=['batch', 3]
    Add(A, B) -> sum
    Sub(A, B) -> diff
    output: name='sum' type='NOTENSOR' shape=None
    output: name='diff' type='NOTENSOR' shape=None

How to validate ONNX outputs against numpy#

Run the exported graph with onnxruntime and compare with the original numpy implementation.

<<<

import numpy as np
import onnxruntime
from yobx.sql import trace_numpy_to_onnx


def normalize_rows(X):
    norms = np.sqrt(np.sum(X ** np.float32(2), axis=1, keepdims=True))
    safe_norms = np.where(norms < np.float32(1e-8), np.float32(1), norms)
    return X / safe_norms


rng = np.random.default_rng(0)
X_sample = rng.standard_normal((8, 4)).astype(np.float32)
X_test = rng.standard_normal((5, 4)).astype(np.float32)

onx = trace_numpy_to_onnx(normalize_rows, X_sample)
sess = onnxruntime.InferenceSession(
    onx.SerializeToString(), providers=["CPUExecutionProvider"]
)

(got,) = sess.run(None, {"X": X_test})
expected = normalize_rows(X_test)
print(f"max|diff|={np.abs(got - expected).max():.2e}")
assert np.allclose(got, expected, atol=1e-5)
print("Outputs match ✓")

>>>

    max|diff|=0.00e+00
    Outputs match ✓