renaming

Author: xadupre

_doc/api/export/control_flow.rst

@@ -0,0 +1,6 @@
+
+onnx_diagnostic.export.control_flow_onnx
+========================================
+
+.. automodule:: onnx_diagnostic.export.control_flow_onnx
+    :members:

_doc/api/export/control_flow_onnx.rst

@@ -6,7 +6,7 @@ onnx_diagnostic.export
     :caption: modules
 
     api
-    control_flow
+    control_flow_onnx
     dynamic_shapes
     onnx_plug
     shape_helper

_doc/api/export/index.rst

@@ -4,7 +4,10 @@
 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_onnx
+from onnx_diagnostic.export.control_flow_onnx 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
 

_unittests/ut_export/test_control_flow.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_onnx`)
+        custom loops (see :func:`onnx_diagnostic.export.control_flow_onnx.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_onnx` or
+    :func:`onnx_diagnostic.export.control_flow_onnx.loop_for_onnx` or
     :class:`onnx_diagnostic.export.onnx_plug.EagerDirectReplacementWithOnnx`.
     """
     if exporter_kwargs and "inline" in exporter_kwargs:
@@ -86,7 +86,7 @@ def to_onnx(
             from experimental_experiment.torch_interpreter import Dispatcher
 
             if use_control_flow_dispatcher:
-                from .control_flow import create_global_dispatcher
+                from .control_flow_onnx import create_global_dispatcher
 
                 control_flow_dispatcher = create_global_dispatcher()
             else:

onnx_diagnostic/export/api.py

@@ -14,7 +14,7 @@
 )
 from torch.fx.experimental.proxy_tensor import ProxyTorchDispatchMode, track_tensor_tree
 from torch.utils._python_dispatch import _get_current_dispatch_mode
-from .control_flow import _loop_for_fn
+from .control_flow_onnx import _loop_for_onnx_fn
 
 
 class SimpleLoopForOp(HigherOrderOperator):
@@ -66,7 +66,7 @@ def simple_loop_for(
         return simple_loop_for_op(n_iter, body_fn, (n_iter, *operands))
 
     if isinstance(n_iter, (bool, int, float)):
-        return _loop_for_fn(body_fn, n_iter, None, *operands)
+        return _loop_for_onnx_fn(body_fn, n_iter, None, *operands)
 
     def _validate_input(n_iter, body_fn, operands):
         assert isinstance(
@@ -127,7 +127,7 @@ def loop_for_op_dense(n_iter, body_fn, operands):
     ), f"Dense implementation operands must be a list of tensors and ints {operands}"
     mode = _get_current_dispatch_mode()
     assert mode is None, "Mode should never be enabled for CPU/CUDA key"
-    return _loop_for_fn(body_fn, n_iter, None, *operands)
+    return _loop_for_onnx_fn(body_fn, n_iter, None, *operands)
 
 
 @simple_loop_for_op.py_impl(ProxyTorchDispatchMode)