extend documentation

Author: xadupre

_doc/status/exported_program_dynamic.rst

@@ -38,6 +38,8 @@ with different options. This steps happens before converting into ONNX.
         print(name)
         print("=" * len(name))
         print()
+        print(f"code: :class:`onnx_diagnostic.torch_export_patches.eval.model_cases.{name}`")
+        print()
         print("forward")
         print("+++++++")
         print()

_doc/status/exporter_dynamic.rst

@@ -3,8 +3,7 @@ Exported ONNX with Dynamic Shapes
 =================================
 
 The following script shows the exported program for many short cases
-and various l-plot-export-with-dynamic-shape to retrieve an ONNX model equivalent
-to the original model.
+to retrieve an ONNX model equivalent to the original model.
 
 .. runpython::
     :showcode:
@@ -40,6 +39,8 @@ to the original model.
         print(name)
         print("=" * len(name))
         print()
+        print(f"code: :class:`onnx_diagnostic.torch_export_patches.eval.model_cases.{name}`")
+        print()
         print("forward")
         print("+++++++")
         print()

onnx_diagnostic/torch_export_patches/eval/model_cases.py

@@ -570,6 +570,34 @@ def forward(self, images, position):
     _dynamic = {"images": {0: DYN, 1: DYN}, "position": {0: DYN}}
 
 
+class ControlFlowWhileDec(torch.nn.Module):
+    def forward(self, ci, a, b):
+        def cond_fn(i, x, y):
+            return i > 0
+
+        def body_fn(i, x, y):
+            return i - 1, x + y, y - x
+
+        return torch._higher_order_ops.while_loop(cond_fn, body_fn, [ci, a, b])
+
+    _inputs = [(torch.tensor(1), torch.randn(2, 3), torch.randn(2, 3))]
+    _dynamic = {}, {0: DYN, 1: DYN}, {0: DYN}
+
+
+class ControlFlowWhileInc(torch.nn.Module):
+    def forward(self, ci, a, b):
+        def cond_fn(i, x, y):
+            return i < x.size(0)
+
+        def body_fn(i, x, y):
+            return i + 1, x + y, y - x
+
+        return torch._higher_order_ops.while_loop(cond_fn, body_fn, [ci, a, b])
+
+    _inputs = [(torch.tensor(1), torch.randn(2, 3), torch.randn(2, 3))]
+    _dynamic = {}, {0: DYN, 1: DYN}, {0: DYN}
+
+
 class SignatureInt1(torch.nn.Module):
     def __init__(self, n_dims: int = 3, n_targets: int = 1):
         super().__init__()