doc

Author: xadupre

onnx_diagnostic/helpers/onnx_helper.py

@@ -46,7 +46,7 @@ def onnx_lighten(
     """
     Creates a model without big initializers but stores statistics
     into dictionaries. The function can be reversed with
-    :func:`onnx_diagnostic.onnx_helper.onnx_unlighten`.
+    :func:`onnx_diagnostic.helpers.onnx_helper.onnx_unlighten`.
     The model is modified inplace.
 
     :param onx: model
@@ -103,7 +103,7 @@ def onnx_unlighten(
 ) -> ModelProto:
     """
     Function fixing the model produced by function
-    :func:`onnx_diagnostic.onnx_helper.onnx_lighten`.
+    :func:`onnx_diagnostic.helpers.onnx_helper.onnx_lighten`.
     The model is modified inplace.
 
     :param onx: model

onnx_diagnostic/torch_onnx/sbs.py

@@ -178,10 +178,9 @@ def run_aligned(
         import torch
         from onnx_diagnostic.reference import (
             # This can be replace by any runtime taking NodeProto as an input.
-            ExtendedReferenceEvaluator as ReferenceEvaluator
+            ExtendedReferenceEvaluator as ReferenceEvaluator,
         )
-        from onnx_diagnostic.torch_interpreter import to_onnx
-        from onnx_diagnostic.torch_interpreter.investigate_helper import run_aligned
+        from onnx_diagnostic.torch_onnx.sbs import run_aligned
 
 
         class Model(torch.nn.Module):
@@ -205,17 +204,17 @@ def post_process(obs):
         ep = torch.export.export(
             Model(), (x,), dynamic_shapes=({0: torch.export.Dim("batch")},)
         )
-        onx = to_onnx(ep)
+        onx = torch.onnx.export(
+            Model(), (x,), dynamic_shapes=({0: torch.export.Dim("batch")},), dynamo=True
+        ).model_proto
         results = list(
             map(
                 post_process,
                 run_aligned(
                     ep,
                     onx,
                     (x,),
-                    check_conversion_cls=dict(
-                        cls=ReferenceEvaluator, atol=1e-5, rtol=1e-5
-                    ),
+                    check_conversion_cls=dict(cls=ReferenceEvaluator, atol=1e-5, rtol=1e-5),
                     verbose=1,
                 ),
             ),