yet-another-onnxruntime-extensions

yet-another-onnxruntime-extensions (yaourt) is an experimental library of ONNX Runtime extensions: custom C++ operators, profiling utilities, and plotting helpers.

Features

• Custom C++ operators (yaourt.ortops) — sparse CPU operators ship as pre-built binaries inside the wheel; fused-kernel CUDA operators require a CUDA-enabled CMake build (see docs/getting_started.rst). Both are registered directly with ONNX Runtime.
• Profiling tools (yaourt.tools) — parse ONNX Runtime JSON profiling files into pandas DataFrames and visualize execution timelines and per-operator breakdowns with matplotlib.
• Plot helpers (yaourt.plot) — benchmark plotting and histogram utilities for model analysis.
• Reference evaluator (yaourt.reference) — a pure-Python ONNX evaluator useful for testing and debugging custom operators without a full ONNX Runtime build.

Installation

pip install yet-another-onnxruntime-extensions

Note: The pre-built wheel includes sparse CPU operators only. To get fused-kernel CUDA operators, install from source with the CUDA toolkit (including nvcc) available on your PATH:

pip install yet-another-onnxruntime-extensions-cuda

See docs/getting_started.rst for full build instructions.

Verify the installation:

import yaourt
print(yaourt.__version__)

Quick Start

Run inference with ONNX Runtime

import numpy as np
import onnxruntime
from yaourt.doc import demo_mlp_model

# Build a small demo MLP model (filename argument is unused)
model = demo_mlp_model("")

# Run inference
sess = onnxruntime.InferenceSession(
    model.SerializeToString(), providers=["CPUExecutionProvider"]
)
x = np.random.randn(3, 10).astype(np.float32)
(output,) = sess.run(None, {"x": x})
print("Output shape:", output.shape)

Load the custom C++ operators

import onnxruntime as ort
from yaourt.ortops import SPARSE_CPU_LIB_PATH

opts = ort.SessionOptions()
opts.register_custom_ops_library(str(SPARSE_CPU_LIB_PATH))

Profile an ONNX Runtime session

from onnxruntime import InferenceSession, SessionOptions
from yaourt.tools.js_profile import js_profile_to_dataframe, plot_ort_profile
import matplotlib.pyplot as plt

opts = SessionOptions()
opts.enable_profiling = True
opts.profile_file_prefix = "/tmp/ort_profile"

sess = InferenceSession(model.SerializeToString(), sess_options=opts,
                        providers=["CPUExecutionProvider"])
# ... run inference ...
profile_file = sess.end_profiling()

df = js_profile_to_dataframe(profile_file, first_it_out=True)
fig, ax = plt.subplots(figsize=(8, 4))
plot_ort_profile(df, ax0=ax, title="Time per operator (µs)")
plt.tight_layout()
plt.show()

Documentation

Full documentation (API reference, examples, getting started guide) is available at: https://xadupre.github.io/docs/yet-another-onnxruntime-extensions/index.html

Contributing

Contributions are welcome! Please read the Getting Started for Developers guide for instructions on how to clone, build, test, and submit changes.

The project uses black for formatting and ruff for linting. Run both before committing:

black . && ruff check .

License

MIT