boundlab.interp.onnx_export#

boundlab.interp.onnx_export(f, args, kwargs={}, input_names=None, output_names=None, optimize=None)[source]#

Export a PyTorch function or module to an ONNX IR model.

Shape arguments are given as lists/tuples of ints — zero-value tensors are constructed internally and only used for tracing.

Parameters:

f (Callable[[...], torch.Tensor] | torch.nn.Module) – A callable or torch.nn.Module to export.
args (tuple[torch.Size | list[int] | torch.Tensor, ...]) – Input shapes, one per positional argument (e.g. ([3, 4],) for a single rank-2 input).
kwargs (dict[str, torch.Size | list[int]]) – Keyword-argument shapes (rarely needed).
input_names (Sequence[str] | None) – Optional names for the ONNX graph inputs.
output_names (Sequence[str] | None) – Optional names for the ONNX graph outputs.

Returns:

An onnx_ir.Model ready for abstract interpretation.

Return type:

onnx_ir.Model

Examples

>>> import torch
>>> from boundlab.interp.onnx import onnx_export
>>> def f(x):
...     return x @ x.T
...
>>> model = onnx_export(f, ([3, 4],))
>>> list(model.graph)[0].op_type
'MatMul'