rename loop_for into loop_for_onnx

Author: xadupre

CHANGELOGS.rst

@@ -19,7 +19,7 @@ Change Logs
 
 * :pr:`303`: fix inputs for summarization, feature extraction tasks
 * :pr:`302`: adds helpers to analyse onnxruntime profiling
-* :pr:`297`: experiment around a higher ops ``loop_for``
+* :pr:`297`: experiment around a higher ops ``loop_for_onnx``
 * :pr:`292`, :pr:`293`, :pr:`294`, :pr:`295`: new patches for Qwen models
 
 0.8.1

_unittests/ut_export/test_control_flow.py

@@ -4,7 +4,7 @@
 from onnxscript import script, FLOAT, INT64
 from onnxscript import opset18 as op
 from onnx_diagnostic.ext_test_case import ExtTestCase, requires_torch, never_test
-from onnx_diagnostic.export.control_flow import enable_code_export_control_flow, loop_for
+from onnx_diagnostic.export.control_flow import enable_code_export_control_flow, loop_for_onnx
 from onnx_diagnostic.export.control_flow_research import simple_loop_for as loop_for_r
 from onnx_diagnostic.export.api import to_onnx
 
@@ -54,7 +54,7 @@ def forward(self, n_iter, x):
                 def body(i, x):
                     return x[: i.item() + 1].unsqueeze(1)
 
-                return loop_for(n_iter, body, (x,))
+                return loop_for_onnx(n_iter, body, (x,))
 
         model = Model()
         n_iter = torch.tensor(4, dtype=torch.int64)
@@ -67,7 +67,7 @@ def body(i, x):
             model, (n_iter, x), dynamic_shapes=({}, ({0: torch.export.Dim.DYNAMIC}))
         )
         self.assertIn(
-            "torch.ops.onnx_higher_ops.loop_for_TestControlFlow_test_loop_one_custom_L_Model_forward_L_body_",
+            "torch.ops.onnx_higher_ops.loop_for_onnx_TestControlFlow_test_loop_one_custom_L_Model_forward_L_body_",
             str(ep),
         )
 
@@ -88,7 +88,7 @@ def forward(self, n_iter, x):
                 def body(i, x):
                     return x[: i.item() + 1].unsqueeze(1)
 
-                return loop_for(n_iter, body, (x,))
+                return loop_for_onnx(n_iter, body, (x,))
 
         model = Model()
         n_iter = torch.tensor(4, dtype=torch.int64)
@@ -125,7 +125,7 @@ def forward(self, n_iter, x):
                 def body(i, x):
                     return x[: i.item() + 1].unsqueeze(1), x[: i.item() + 1].unsqueeze(1) + 1
 
-                res = loop_for(n_iter, body, (x,))
+                res = loop_for_onnx(n_iter, body, (x,))
                 return res[0] + res[1]
 
         model = Model()
@@ -139,7 +139,7 @@ def body(i, x):
             model, (n_iter, x), dynamic_shapes=({}, ({0: torch.export.Dim.DYNAMIC}))
         )
         self.assertIn(
-            "torch.ops.onnx_higher_ops.loop_for_TestControlFlow_test_loop_two_custom_L_Model_forward_L_body_",
+            "torch.ops.onnx_higher_ops.loop_for_onnx_TestControlFlow_test_loop_two_custom_L_Model_forward_L_body_",
             str(ep),
         )
 
@@ -160,7 +160,7 @@ def forward(self, n_iter, x):
                 def body(i, x):
                     return x[: i.item() + 1].unsqueeze(1), x[: i.item() + 1].unsqueeze(0) + 1
 
-                res = loop_for(n_iter, body, (x,), reduction_dim=[0, 1])
+                res = loop_for_onnx(n_iter, body, (x,), reduction_dim=[0, 1])
                 return res[0] + res[1].T
 
         model = Model()
@@ -174,7 +174,7 @@ def body(i, x):
             model, (n_iter, x), dynamic_shapes=({}, ({0: torch.export.Dim.DYNAMIC}))
         )
         self.assertIn(
-            "torch.ops.onnx_higher_ops.loop_for_TestControlFlow_test_loop_two_custom_reduction_dim_L_Model_forward_L_body_",
+            "torch.ops.onnx_higher_ops.loop_for_onnx_TestControlFlow_test_loop_two_custom_reduction_dim_L_Model_forward_L_body_",
             str(ep),
         )
 

onnx_diagnostic/export/api.py

@@ -43,7 +43,7 @@ def to_onnx(
     :param save_ep: saves the exported program
     :param optimize: optimizes the model
     :param use_control_flow_dispatcher: use the dispatcher created to supported
-        custom loops (see :func:`onnx_diagnostic.export.control_flow.loop_for`)
+        custom loops (see :func:`onnx_diagnostic.export.control_flow.loop_for_onnx`)
     :param onnx_plugs: the code was modified to replace some parts with onnx translation
     :param inline: inline local functions
     :return: the output of the selected exporter, usually a structure including
@@ -62,7 +62,7 @@ def to_onnx(
         )
 
     Some examples using control flows are available in
-    :func:`onnx_diagnostic.export.control_flow.loop_for` or
+    :func:`onnx_diagnostic.export.control_flow.loop_for_onnx` or
     :class:`onnx_diagnostic.export.onnx_plug.EagerDirectReplacementWithOnnx`.
     """
     if exporter_kwargs and "inline" in exporter_kwargs:

onnx_diagnostic/export/control_flow.py

@@ -55,7 +55,7 @@ def is_exporting() -> bool:
     return _TEST_EXPORT or torch.compiler.is_exporting() or torch.compiler.is_compiling()
 
 
-def _loop_for_fn(n_iter, body_fn, reduction_dim, args):
+def _loop_for_onnx_fn(n_iter, body_fn, reduction_dim, args):
     """
     Python implementation of the loop.
 
@@ -103,7 +103,7 @@ def _loop_for_fn(n_iter, body_fn, reduction_dim, args):
     return tuple(final) if len(final) > 1 else final[0]
 
 
-def make_custom_loop_for(
+def make_custom_loop_for_onnx(
     n_iter: torch.Tensor,
     body_fn: Callable,
     reduction_dim: Optional[Sequence[int]],
@@ -139,7 +139,7 @@ def make_custom_loop_for(
         .replace("<lambda>", "l")
         .replace(".", "_")
     )
-    name = f"loop_for_{full_name}_{srank}_{sred}"
+    name = f"loop_for_onnx_{full_name}_{srank}_{sred}"
     if name in _REGISTERED_SCHEMA:
         return name, _REGISTERED_SCHEMA[name][0]
     sig = inspect.signature(body_fn)
@@ -197,10 +197,10 @@ def convert_custom_loop_into_onnx(
     *args: str,
     body_callable: Callable[..., onnx.ModelProto],
     reduction_dim: Optional[Sequence[int]] = None,
-    name: str = "loop_for",
+    name: str = "loop_for_onnx",
 ) -> Union[str, List[str]]:
     """
-    Converts a custom op ``higher_ops::loop_for...`` into e sequence of node.
+    Converts a custom op ``higher_ops::loop_for_onnx...`` into e sequence of node.
 
     :param g: GreaphBuilder
     :param sts: if not defined, torch does not know the output shapes
@@ -267,7 +267,7 @@ def convert_custom_loop_into_onnx(
 
     sequences = [g.op.SequenceEmpty() for _ in outputs]
 
-    outloop = [g.unique_name(f"loop_for{i}") for i in range(len(sequences))]
+    outloop = [g.unique_name(f"loop_for_onnx{i}") for i in range(len(sequences))]
 
     for i, s in enumerate(sequences):
         g.set_sequence(s, graph.output[i].type.tensor_type.elem_type)
@@ -321,7 +321,7 @@ def convert_into_onnx(
     return container.model_proto
 
 
-def loop_for(
+def loop_for_onnx(
     n_iter: Union[torch.SymInt, torch.Tensor],
     body_fn: Callable[..., Tuple[torch.Tensor]],
     args: Sequence[torch.Tensor],
@@ -352,15 +352,15 @@ def loop_for(
         import torch
         import onnxruntime
         from onnx_diagnostic.export.api import to_onnx
-        from onnx_diagnostic.export.control_flow import loop_for
+        from onnx_diagnostic.export.control_flow import loop_for_onnx
 
 
         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,))
+                return loop_for_onnx(n_iter, body, (x,))
 
 
         model = Model()
@@ -398,15 +398,15 @@ def body(i, x):
         import torch
         import onnxruntime
         from onnx_diagnostic.export.api import to_onnx
-        from onnx_diagnostic.export.control_flow import loop_for
+        from onnx_diagnostic.export.control_flow import loop_for_onnx
 
 
         class Model(torch.nn.Module):
             def forward(self, n_iter, x):
                 def body(i, x):
                     return x[: i.item() + 1].unsqueeze(1), x[: i.item() + 1].unsqueeze(1) + 1
 
-                two = loop_for(n_iter, body, (x,))
+                two = loop_for_onnx(n_iter, body, (x,))
                 return two[0] + two[1]
 
 
@@ -445,15 +445,15 @@ def body(i, x):
         import torch
         import onnxruntime
         from onnx_diagnostic.export.api import to_onnx
-        from onnx_diagnostic.export.control_flow import loop_for
+        from onnx_diagnostic.export.control_flow import loop_for_onnx
 
 
         class Model(torch.nn.Module):
             def forward(self, n_iter, x):
                 def body(i, x):
                     return x[: i.item() + 1].unsqueeze(1), x[: i.item() + 1].unsqueeze(0) + 1
 
-                two = loop_for(n_iter, body, (x,), reduction_dim=[0, 1])
+                two = loop_for_onnx(n_iter, body, (x,), reduction_dim=[0, 1])
                 return two[0] + two[1].T
 
 
@@ -501,7 +501,7 @@ def body(i, x):
             body_mutated_inputs,
             body_outputs,
         ) = check_input_alias_and_mutation_return_outputs(body_gm)
-        name, _custom_ops = make_custom_loop_for(
+        name, _custom_ops = make_custom_loop_for_onnx(
             n_iter,
             body_fn,
             reduction_dim,
@@ -513,4 +513,4 @@ def body(i, x):
         fct = getattr(torch.ops.onnx_higher_ops, name)
         return fct(n_iter, *args)
 
-    return _loop_for_fn(n_iter, body_fn, reduction_dim, args)
+    return _loop_for_onnx_fn(n_iter, body_fn, reduction_dim, args)