yobx.to_onnx#

yobx.to_onnx() is the single entry point for converting any supported model to ONNX format. It inspects the type of the model argument at runtime and automatically delegates to the appropriate backend-specific converter, forwarding all extra keyword arguments verbatim.

The function always returns an ExportArtifact regardless of which backend was selected — see to_onnx and ExportArtifact for a full description of that container.

This page documents the user-facing API and implementation details of the dispatcher in yobx.convert.

Dispatch rules#

yobx.to_onnx() inspects model in this order and calls the first matching backend:

Model type	Backend function	Reference
`torch.nn.Module` or `torch.fx.GraphModule`	`yobx.torch.to_onnx()`	Torch Export to ONNX
`sklearn.base.BaseEstimator`	`yobx.sklearn.to_onnx()`	scikit-learn Export to ONNX
`tf.Module` (inc. Keras)	`yobx.tensorflow.to_onnx()`	TensorFlow / JAX Export to ONNX
`bytes` or `*.tflite` path	`yobx.litert.to_onnx()`	LiteRT / TFLite Export to ONNX
SQL `str`, `callable()`, or `polars.LazyFrame`	`yobx.sql.to_onnx()`	SQL-to-ONNX Converter

to_onnx(model, args, ...)
     │
     ├── torch.nn.Module / torch.fx.GraphModule ──► yobx.torch.to_onnx
     │
     ├── sklearn.base.BaseEstimator ──────────────► yobx.sklearn.to_onnx
     │
     ├── tensorflow.Module ───────────────────────► yobx.tensorflow.to_onnx
     │
     ├── bytes / *.tflite path ───────────────────► yobx.litert.to_onnx
     │
     └── str / callable / polars.LazyFrame ──────► yobx.sql.to_onnx

If the model type matches none of the above, a TypeError is raised listing all supported types.

Common parameters#

All backends accept the following keyword arguments. Backend-specific parameters can be passed as extra **kwargs and are forwarded verbatim to the selected converter.

Parameter	Default	Description
`args`	`None`	Input arguments forwarded to the selected converter. For torch: a tuple of `torch.Tensor` objects. For sklearn / tensorflow / litert: a tuple of `numpy.ndarray` objects. For sql: a `{column: dtype}` mapping or a `numpy.ndarray` for numpy-function tracing.
`input_names`	`None`	Optional list of names for the ONNX graph input tensors. When omitted, names are derived automatically by each backend.
`dynamic_shapes`	`None`	Declares which tensor dimensions are symbolic (variable-length). See Dynamic shapes below for the backend-specific formats.
`target_opset`	`21`	ONNX opset version to target. Either an integer for the default domain, or a `Dict[str, int]` mapping domain names to opset versions (e.g. `{"": 21, "com.microsoft": 1}`).
`verbose`	`0`	Verbosity level (0 = silent).
`large_model`	`False`	When `True`, initializers are stored as external data rather than embedded in the proto. Required for models exceeding the 2 GB protobuf limit.
`external_threshold`	`1024`	When `large_model=True`, tensors whose element count exceeds this threshold are stored externally.
`filename`	`None`	If set, saves the exported model to this path. Not supported by the LiteRT backend (silently ignored).
`return_optimize_report`	`False`	When `True`, the returned artifact’s `report` attribute is populated with per-pattern optimization statistics.

Dynamic shapes#

The meaning of None and the accepted format for dynamic_shapes differs by backend:

PyTorch follows torch.export.export() conventions. Pass a dict mapping input names to per-axis torch.export.Dim objects, or a tuple with one entry per positional input. None means no dynamic dimensions (all shapes are fixed):

import torch
from yobx import to_onnx

batch = torch.export.Dim("batch", min=1, max=256)
artifact = to_onnx(model, (x,), dynamic_shapes={"x": {0: batch}})

scikit-learn / tensorflow / LiteRT use a tuple of {axis: dim_name} dicts, one per input. None treats axis 0 as dynamic (batch dimension) for every input; all other axes remain static:

import numpy as np
from sklearn.linear_model import LinearRegression
from yobx import to_onnx

# None → axis 0 is automatically dynamic
artifact = to_onnx(reg, (X,))

# explicit dynamic batch for input 0, static for input 1
artifact = to_onnx(reg, (X,), dynamic_shapes=({0: "batch"},))

SQL / numpy callable uses the same {axis: dim_name} tuple format as scikit-learn. None lets the backend apply its own default (typically axis 0 dynamic). Ignored for SQL strings and DataFrame-tracing callables.

Return value#

Every call to yobx.to_onnx() returns an ExportArtifact instance. See to_onnx and ExportArtifact for the complete API, including how to:

access the ONNX proto via artifact.proto or artifact.get_proto()
save the model to disk with artifact.save("model.onnx")
reload from disk with ExportArtifact.load("model.onnx")
inspect per-pattern optimization statistics via artifact.report

Examples#

scikit-learn — linear regression#

import numpy as np
from sklearn.linear_model import LinearRegression
from yobx import to_onnx

X = np.random.randn(20, 4).astype(np.float32)
y = X[:, 0] + X[:, 1]
reg = LinearRegression().fit(X, y)

# axis 0 is treated as dynamic (batch) by default
artifact = to_onnx(reg, (X,))
print(artifact)

scikit-learn — explicit dynamic batch dimension#

import numpy as np
from sklearn.linear_model import LinearRegression
from yobx import to_onnx

X = np.random.randn(20, 4).astype(np.float32)
y = X[:, 0] + X[:, 1]
reg = LinearRegression().fit(X, y)

# Mark axis 0 of the first input as dynamic:
artifact = to_onnx(reg, (X,), dynamic_shapes=({0: "batch"},))

PyTorch — dynamic batch dimension#

import torch
from yobx import to_onnx

class Neuron(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(4, 2)
    def forward(self, x):
        return torch.relu(self.linear(x))

model = Neuron()
x = torch.randn(3, 4)
batch = torch.export.Dim("batch", min=1, max=256)
artifact = to_onnx(model, (x,), dynamic_shapes={"x": {0: batch}})
artifact.save("neuron.onnx")

TensorFlow / Keras — simple model#

import numpy as np
import tensorflow as tf
from yobx import to_onnx

model = tf.keras.Sequential([
    tf.keras.layers.Dense(4, activation="relu"),
    tf.keras.layers.Dense(1),
])
X = np.random.randn(10, 3).astype(np.float32)
model(X)   # build the model
artifact = to_onnx(model, (X,))

SQL query#

import numpy as np
from yobx import to_onnx

artifact = to_onnx(
    "SELECT a + b AS total FROM t WHERE a > 0",
    {"a": np.float32, "b": np.float32},
)

LiteRT / TFLite — from file#

from yobx import to_onnx

artifact = to_onnx("model.tflite", ())

LiteRT / TFLite — from raw bytes#

from yobx import to_onnx

with open("model.tflite", "rb") as f:
    flatbuffer = f.read()

artifact = to_onnx(flatbuffer, ())

Saving and running the exported model#

Once you have an ExportArtifact, you can serialize it and run it with any ONNX-compatible runtime:

import numpy as np
import onnxruntime
from sklearn.linear_model import LinearRegression
from yobx import to_onnx

X = np.random.randn(20, 4).astype(np.float32)
y = X[:, 0] + X[:, 1]
reg = LinearRegression().fit(X, y)

artifact = to_onnx(reg, (X,))
artifact.save("reg.onnx")

sess = onnxruntime.InferenceSession("reg.onnx")
(predictions,) = sess.run(None, {"X": X[:5]})

OnnxRuntime operator fusions#

Passing "com.microsoft": 1 in target_opset enables operator fusions (fused attention, layer normalization, …) that are specific to onnxruntime:

from yobx import to_onnx

artifact = to_onnx(
    model,
    (x,),
    target_opset={"": 21, "com.microsoft": 1},
)

Implementation#

The dispatcher is implemented in yobx/convert.py and exported from the top-level yobx package via:

from .convert import DEFAULT_TARGET_OPSET, to_onnx

It resolves the backend at call time (not at import time) so that optional dependencies such as torch, scikit-learn, or tensorflow do not need to be installed for unrelated backends to work. Each backend module is imported inside the matching if branch using a local import.

Backend availability is checked with the helpers has_torch(), has_sklearn(), has_tensorflow(), and has_litert() from yobx.ext_test_case, which perform a lightweight importlib.util.find_spec probe without importing the package itself.

Dispatch order#

The checks run in a fixed priority order:

torch — torch.nn.Module or torch.fx.GraphModule
sklearn — sklearn.base.BaseEstimator
tensorflow — tf.Module
litert — bytes raw flatbuffer, or a path / string that ends with ".tflite" and exists on disk
sql — any str, os.PathLike, or callable() (including polars.LazyFrame, whose type is duck-typed via type(model).__module__)

If none of the checks match, a TypeError is raised with a descriptive message.

Common arguments dict#

Before branching, the dispatcher assembles a common dict that collects all shared keyword arguments:

common = dict(
    input_names=input_names,
    dynamic_shapes=dynamic_shapes,
    target_opset=target_opset,
    verbose=verbose,
    large_model=large_model,
    external_threshold=external_threshold,
    filename=filename,
    return_optimize_report=return_optimize_report,
)

Each backend receives **common, **kwargs so that the caller’s extra keyword arguments are always passed through without being modified by the dispatcher.

LiteRT special case#

The LiteRT backend does not support the filename parameter. The dispatcher therefore builds a separate litert_common dict with filename removed before forwarding to yobx.litert.to_onnx():

litert_common = {k: v for k, v in common.items() if k != "filename"}

API reference#

yobx.to_onnx(model: Any, args: Any = None, input_names: Sequence[str] | None = None, dynamic_shapes: Any | None = None, target_opset: int | Dict[str, int] = 21, verbose: int = 0, large_model: bool = False, external_threshold: int = 1024, filename: str | None = None, return_optimize_report: bool = False, **kwargs: Any) → Any[source]#

Convert any supported model type to ONNX.

This is the unified top-level entry point. It inspects model and args to decide which backend converter to call:

torch — when model is a torch.nn.Module or torch.fx.GraphModule. Delegates to yobx.torch.to_onnx(). Any extra kwargs are forwarded verbatim; see that function for the full list of accepted parameters (as_function, options, dispatcher, dynamic_shapes, export_options, function_options, …).
scikit-learn — when model is a sklearn.base.BaseEstimator. Delegates to yobx.sklearn.to_onnx(). Extra kwargs include builder_cls, extra_converters, function_options, convert_options, …
TensorFlow / Keras — when model is a tf.Module (including Keras models and layers). Delegates to yobx.tensorflow.to_onnx(). Extra kwargs include builder_cls, extra_converters, …
LiteRT / TFLite — when model is bytes (raw flatbuffer) or a str / os.PathLike whose path ends with ".tflite". Delegates to yobx.litert.to_onnx(). Extra kwargs include builder_cls, extra_converters, subgraph_index, … Note: filename is not supported by this backend and is silently ignored.
SQL / DataFrame / NumPy callable — when model is a plain str (SQL query), a Python callable(), or a polars.LazyFrame. Delegates to yobx.sql.to_onnx(). Extra kwargs include custom_functions, builder_cls, …

Parameters:

model – the model to convert; see dispatch rules above.
args – input arguments passed to the selected converter. For torch this is a sequence of torch.Tensor objects. For sklearn / tensorflow / litert this is a tuple of numpy.ndarray objects. For SQL / DataFrame this is a {column: dtype} mapping or a numpy.ndarray (numpy-function tracing).
input_names – optional list of names for the ONNX graph input tensors. Supported by all backends.
dynamic_shapes –
optional specification of dynamic (symbolic) dimensions. The exact format and default behaviour differ by backend:
- torch — follows torch.export.export() conventions: a dict mapping input names to per-axis torch.export.Dim objects, or a tuple with one entry per input. None means no dynamic dimensions (all shapes are fixed).
- sklearn / tensorflow / litert — a tuple of {axis: dim_name} dicts, one per input. None means axis 0 is treated as dynamic (batch dimension) for every input, while all other axes remain static.
- sql (numpy callable) — same {axis: dim_name} tuple format. None lets the backend apply its own default (typically axis 0 dynamic). Ignored for SQL strings and DataFrame-tracing callables.
target_opset – ONNX opset version to target. Either an integer for the default domain (""), or a Dict[str, int] mapping domain names to opset versions (e.g. {"": 21, "ai.onnx.ml": 5}). Defaults to DEFAULT_TARGET_OPSET.
verbose – verbosity level (0 = silent). Supported by all backends.
large_model – if True the returned ExportArtifact stores tensors as external data rather than embedding them inside the proto. Supported by all backends.
external_threshold – when large_model is True, tensors whose element count exceeds this threshold are stored externally. Supported by all backends.
filename – if set, the exported ONNX model is saved to this path. Supported by torch, sklearn, tensorflow, and sql backends. Not supported by the LiteRT backend (ignored when model is a .tflite file or raw flatbuffer bytes).
return_optimize_report – if True, the returned ExportArtifact has its report attribute populated with per-pattern optimization statistics. Supported by all backends.
kwargs – additional backend-specific keyword arguments forwarded verbatim to the selected converter. See the backend-specific to_onnx functions listed above for their full parameter lists.

Returns:

ExportArtifact wrapping the exported ONNX proto together with an ExportReport.

Example — scikit-learn (batch dimension fixed):

import numpy as np
from sklearn.linear_model import LinearRegression
from yobx import to_onnx

X = np.random.randn(20, 4).astype(np.float32)
y = X[:, 0] + X[:, 1]
reg = LinearRegression().fit(X, y)
artifact = to_onnx(reg, (X,))

Example — scikit-learn (explicit dynamic batch dimension):

import numpy as np
from sklearn.linear_model import LinearRegression
from yobx import to_onnx

X = np.random.randn(20, 4).astype(np.float32)
y = X[:, 0] + X[:, 1]
reg = LinearRegression().fit(X, y)
# Mark axis 0 (batch) as dynamic for input 0:
artifact = to_onnx(reg, (X,), dynamic_shapes=({0: "batch"},))

Example — PyTorch (dynamic batch dimension via torch.export.Dim):

import torch
from yobx import to_onnx

class Neuron(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(4, 2)
    def forward(self, x):
        return torch.relu(self.linear(x))

model = Neuron()
x = torch.randn(3, 4)
batch = torch.export.Dim("batch", min=1, max=256)
artifact = to_onnx(model, (x,), dynamic_shapes={"x": {0: batch}})

Example — SQL:

import numpy as np
from yobx import to_onnx

artifact = to_onnx(
    "SELECT a + b AS total FROM t WHERE a > 0",
    {"a": np.float32, "b": np.float32},
)