yobx.sql.trace_numpy_to_onnx#

yobx.sql.trace_numpy_to_onnx(func: Callable, *inputs: ndarray, input_names: Sequence[str] | None = None, output_names: Sequence[str] | None = None, target_opset: int | Dict[str, int] = 21, batch_dim: str = 'batch', dynamic_shapes: Tuple[Dict[int, str], ...] | None = None, large_model: bool = False, external_threshold: int = 1024, return_optimize_report: bool = False) → ExportArtifact[source]#

Trace a numpy function and return the equivalent ONNX model.

This is the high-level entry point. Internally it creates a fresh GraphBuilder, registers the graph inputs derived from the inputs sample arrays, delegates the actual tracing to trace_numpy_function(), registers the graph outputs, and exports to onnx.ModelProto.

Parameters:

  • func – a Python callable that accepts one or more numpy arrays and returns a numpy array (or a tuple/list of arrays). The function may use numpy ufuncs, arithmetic operators, reductions, and shape manipulations; see yobx.xtracing.numpy_array for the full list of supported operations.

  • inputs – sample numpy arrays used to determine the element type and shape of each input. Only dtype and shape are used; the actual values are ignored.

  • input_names – optional list of tensor names for the ONNX graph inputs. Defaults to ["X"] for a single input or ["X0", "X1", …] for multiple inputs.

  • output_names – optional list of tensor names for the ONNX graph outputs. Defaults to ["output_0"]. For functions that return multiple arrays supply the correct number of names, e.g. ["out_a", "out_b"].

  • target_opset – ONNX opset version. Can be an integer (default domain only) or a dictionary mapping domain names to opset versions.

  • batch_dim – name of the dynamic first dimension (default: "batch"). Used when dynamic_shapes is not provided. Change this when the default name conflicts with another symbolic dimension in the same graph.

  • dynamic_shapes – optional per-input axis-to-dimension-name mappings. Each element is a {axis: dim_name} dict for the corresponding input. When provided, these override the default batch_dim behaviour for each input.

  • large_model – if True the returned ExportArtifact has its container attribute set to an ExtendedModelContainer

  • external_threshold – if large_model is True, every tensor whose element count exceeds this threshold is stored as external data

  • return_optimize_report – if True, the returned ExportArtifact has its report attribute populated with per-pattern optimization statistics

Returns:

an ExportArtifact representing the traced function.

Example:

import numpy as np
from yobx.sql import trace_numpy_to_onnx

def my_func(X):
    return np.sqrt(np.abs(X) + 1)

X = np.random.randn(4, 3).astype(np.float32)
onx = trace_numpy_to_onnx(my_func, X)
On this page

This Page