yobx.xtracing#

Lightweight mechanism for tracing numpy functions and exporting them to ONNX. See Numpy-Tracing and FunctionTransformer for a full walkthrough.

modules

trace_numpy_function#

yobx.xtracing.trace_numpy_function(g: GraphBuilderExtendedProtocol, sts: Dict, outputs: List[str], func: Callable, inputs: List[str], name: str = 'trace', kw_args: Dict[str, Any] | None = None) → str[source]#

Trace a numpy function by wrapping named tensors in g as NumpyArray proxies, then recording all numpy operations as ONNX nodes in g.

This function follows the same API convention as other converters in this package: it takes an existing GraphBuilder, the scikit-learn shape/type dictionary sts, the desired output tensor names, and then the callable and input tensor names.

Parameters:

  • g – the graph builder to add nodes to

  • sts – shapes defined by scikit-learn (may be empty; forwarded to g.set_type_shape_unary_op when non-empty)

  • outputs – desired output tensor names (one per output returned by func)

  • func – the numpy function to trace; must accept len(inputs) positional arguments and return a NumpyArray or a tuple/list of NumpyArray objects

  • inputs – existing input tensor names already present in g

  • name – node name prefix used when emitting Identity rename nodes

  • kw_args – optional keyword arguments forwarded to func

Returns:

the first output tensor name (outputs[0])

Example:

<<<

from yobx.helpers.onnx_helper import pretty_onnx
from yobx.xbuilder import GraphBuilder
from yobx.xtracing import trace_numpy_function
import numpy as np
from onnx import TensorProto

g = GraphBuilder({"": 21, "ai.onnx.ml": 1})
g.make_tensor_input("X", TensorProto.FLOAT, ("batch", 3))


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


trace_numpy_function(g, {}, ["output_0"], my_func, ["X"])
g.make_tensor_output("output_0", indexed=False, allow_untyped_output=True)
art = g.to_onnx()
print(pretty_onnx(art))

>>>

    opset: domain='' version=21
    opset: domain='ai.onnx.ml' version=1
    input: name='X' type=dtype('float32') shape=['batch', 3]
    init: name='init1_s_' type=float32 shape=() -- array([1.], dtype=float32)-- Opset.make_node.1/Small
    Abs(X) -> _onx_abs_X
      Add(_onx_abs_X, init1_s_) -> _onx_add_abs_X
        Sqrt(_onx_add_abs_X) -> output_0
    output: name='output_0' type='NOTENSOR' shape=None

trace_numpy_to_onnx#

yobx.xtracing.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') → 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"). Change this when the default name conflicts with another symbolic dimension in the same graph.

Returns:

an ExportArtifact representing the traced function.

Example:

import numpy as np
from yobx.xtracing 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)