fix a few things

Author: xadupre

_unittests/ut_export/test_control_flow.py

@@ -78,6 +78,40 @@ def body(i, x):
         self.dump_onnx("test_loop_one_custom.onnx", onx)
         self.assert_onnx_disc("test_loop_one_custom", onx, model, (n_iter, x))
 
+    @requires_torch("2.9.99")
+    def test_loop_one_custom_different_opset(self):
+        class Model(torch.nn.Module):
+            def forward(self, n_iter, x):
+                def body(i, x):
+                    return x[: i.item() + 1].unsqueeze(1)
+
+                return loop_for(n_iter, body, (x,))
+
+        model = Model()
+        n_iter = torch.tensor(4, dtype=torch.int64)
+        x = torch.arange(10, dtype=torch.float32)
+        expected = torch.tensor([0, 0, 1, 0, 1, 2, 0, 1, 2, 3], dtype=x.dtype).unsqueeze(1)
+        got = model(n_iter, x)
+        self.assertEqualArray(expected, got)
+
+        ep = torch.export.export(
+            model, (n_iter, x), dynamic_shapes=({}, ({0: torch.export.Dim.DYNAMIC}))
+        )
+        self.assertIn("torch.ops.onnx_higher_ops.loop_for_body_", str(ep))
+
+        onx = to_onnx(
+            model,
+            (n_iter, x),
+            dynamic_shapes=({}, ({0: torch.export.Dim.DYNAMIC})),
+            exporter="custom",
+            use_control_flow_dispatcher=True,
+            target_opset=22,
+        ).model_proto
+        opsets = {d.domain: d.version for d in onx.opset_import}
+        self.assertEqual(opsets[""], 22)
+        self.dump_onnx("test_loop_one_custom.onnx", onx)
+        self.assert_onnx_disc("test_loop_one_custom", onx, model, (n_iter, x))
+
     @requires_torch("2.9.99")
     def test_loop_two_custom(self):
         class Model(torch.nn.Module):

onnx_diagnostic/export/api.py

@@ -66,6 +66,12 @@ def to_onnx(
 
             dispatcher = create_global_dispatcher()
 
+        options = None
+        if exporter_kwargs is not None:
+            options = exporter_kwargs.pop("options", None)
+        if options is None:
+            options = OptimizationOptions(patterns="default+onnxruntime")
+
         return _to_onnx(
             mod,
             args=args,
@@ -79,7 +85,7 @@ def to_onnx(
             large_model=True,
             output_dynamic_shapes=output_dynamic_shapes,
             export_options=ExportOptions(save_ep=save_ep),
-            options=OptimizationOptions(patterns="default+onnxruntime"),
+            options=options,
             **(exporter_kwargs or {}),
             dispatcher=dispatcher if use_control_flow_dispatcher else None,
         )

onnx_diagnostic/export/control_flow.py

@@ -148,14 +148,31 @@ def make_custom_loop_for(
     custom_def._abstract_fn = lambda *_args, _o=body_outputs: (
         tuple([torch.empty_like(s) for s in _o]) if len(_o) > 1 else torch.empty_like(_o[0])
     )
-    onx = convert_into_onnx(body_gm, args)
+
+    def _make_onx(
+        body_gm=body_gm, args=args, target_opset=None, verbose=0, exporter_kwargs=None
+    ):
+        return convert_into_onnx(
+            body_gm,
+            args,
+            exporter_kwargs=exporter_kwargs,
+            target_opset=target_opset,
+            verbose=verbose,
+        )
+
     to_register = (
         custom_def,
-        onx,
+        _make_onx,
         (
-            lambda g, sts, outputs, *args, body=onx, reduction_dim=reduction_dim, name=name: (
+            lambda g, sts, outputs, *args, bc=_make_onx, rd=reduction_dim, name=name: (
                 convert_custom_loop_into_onnx(
-                    g, sts, outputs, *args, body=body, reduction_dim=reduction_dim, name=name
+                    g,
+                    sts,
+                    outputs,
+                    *args,
+                    body_callable=bc,
+                    reduction_dim=rd,
+                    name=name,
                 )
             )
         ),
@@ -173,7 +190,7 @@ def convert_custom_loop_into_onnx(
     sts: Dict[str, Any],
     outputs: List[str],
     *args: str,
-    body: onnx.GraphProto,
+    body_callable: Callable[..., onnx.ModelProto],
     reduction_dim: Optional[Sequence[int]] = None,
     name: str = "loop_for",
 ) -> Union[str, List[str]]:
@@ -190,6 +207,14 @@ def convert_custom_loop_into_onnx(
     :param name: to give the onnx nodes a name
     :return: output names
     """
+    assert body_callable is not None, "body_callable cannot be None"
+    # This should be part of a public API.
+    body = body_callable(
+        target_opset=g.main_opset,
+        verbose=g.verbose,
+        exporter_kwargs={"options": g.optimization_options},
+    )
+
     graph = body.graph if isinstance(body, onnx.ModelProto) else body
     assert isinstance(
         graph, onnx.GraphProto
@@ -261,19 +286,33 @@ def convert_custom_loop_into_onnx(
 
 
 def convert_into_onnx(
-    body_gm: torch.fx.GraphModule, args: Sequence[torch.Tensor]
+    body_gm: torch.fx.GraphModule,
+    args: Sequence[torch.Tensor],
+    target_opset: Optional[int] = None,
+    verbose: int = 0,
+    exporter_kwargs: Optional[Dict[str, Any]] = None,
 ) -> onnx.ModelProto:
     """
     Converts a torch.fx.GraphModule into ONNX.
     It returns a ModelProto.
 
     :param body_gm: a torch.fx.GraphModule
     :param args: arguments known at export time
+    :param target_opset: targetted opset
+    :param verbose: verbosity level
+    :param exporter_kwargs: additional exporter arguments
     :return: a ModelProto
     """
     # This does not work with onnx-dynamo.
     # opset still needs to be defined
-    container = to_onnx(body_gm, args, exporter="custom")
+    container = to_onnx(
+        body_gm,
+        args,
+        exporter="custom",
+        exporter_kwargs=exporter_kwargs,
+        target_opset=target_opset,
+        verbose=verbose,
+    )
     return container.model_proto