onnx-diagnostic: investigate onnx models

The main feature is about patches: it helps exporting pytorch models into ONNX, mostly designed for LLMs using dynamic caches. Sources available at github/onnx-diagnostic.

with bypass_export_some_errors(patch_transformers=True) as f:
    ep = torch.export.export(model, args, kwargs=kwargs, dynamic_shapes=dynamic_shapes)
    # ...

It also implements tools to investigate, validate exported models (ExportedProgramm, ONNXProgram, …). onnx_diagnostic.torch_export_patches.bypass_export_some_errors().

Contents

• API of onnx_diagnostic
• Command Lines
• Examples Gallery
• Common Export Issues

More

• Change Logs
• License

Getting started

git clone https://github.com/sdpython/onnx-diagnostic.git
cd onnx-diagnostic
pip install -e .

or

pip install onnx-diagnostic

Enlightening Examples

Torch Export

• Export a model with a control flow (If)

• Use DYNAMIC or AUTO when exporting if dynamic shapes has constraints

• Export with DynamicCache and dynamic shapes

• Find and fix an export issue due to dynamic shapes

• Steel method forward to guess the dynamic shapes (with Tiny-LLM)

• Export Tiny-LLM with patches

• Test the export on untrained models

Investigate ONNX models

• Intermediate results with (ONNX) ReferenceEvaluator

• Intermediate results with onnxruntime

• Find where a model is failing by running submodels

Some Usefuls Tools

import torch
from onnx_diagnostic.helpers import string_type

inputs = (
    torch.rand((3, 4), dtype=torch.float16),
    [
        torch.rand((5, 6), dtype=torch.float16),
        torch.rand((5, 6, 7), dtype=torch.float16),
    ]
)

# with shapes
print(string_type(inputs, with_shape=True))

>>> (T10s3x4,#2[T10s5x6,T10s5x6x7])

import onnx
from onnx_diagnostic.helpers import onnx_dtype_name

itype = onnx.TensorProto.BFLOAT16
print(onnx_dtype_name(itype))
print(onnx_dtype_name(7))

>>> BFLOAT16
>>> INT64

onnx_diagnostic.helpers.max_diff(), …

The documentation was updated on:

    2025-04-25 15:02:36.019013

With the following versions:

    numpy: 2.2.5
    ml_dtypes: 0.5.1
    sklearn: 1.6.1
    onnx: 1.19.0
    onnxruntime: 1.22.0+cu126
    onnxscript: 0.3.0.dev20250301
    torch: 2.8.0.dev20250423+cu126
    transformers: 4.52.0.dev0
    timm: 1.0.14
    has_onnxruntime_training: True

Size of the package:

                              lines   chars
    ext dir                                
    .py                        1396   39640
        export                  926   29144
        helpers                3841  111389
        reference              1337   38371
        tasks                  1119   37024
        torch_export_patches   1502   53340
        torch_models           5203  272134
        torch_onnx              349   11207

Older versions

• 0.4.2

• 0.4.1

• 0.4.0

• 0.3.0

• 0.2.2

• 0.2.1

• 0.2.0

• 0.1.0