onnx-diagnostic 0.8.6__py3-none-any.whl → 0.8.8__py3-none-any.whl

onnx_diagnostic/ci_models/export_qwen25_vl.py

@@ -8,10 +8,10 @@ Requirements
 ::
 
     git+https://github.com/sdpython/experimental-experiment.git  # optional
-    huggingface_hub>=1.2.1
+    huggingface_hub
     onnx-diagnostic>=0.8.6
     onnxruntime>=1.23
-    torch>=2.9  # weekly is better
+    torch>=2.10  # weekly is better
     tqdm
     transformers>=4.57
 
@@ -59,6 +59,7 @@ It is possible to overwrite this by by setting environment variable
 import os
 import sys
 import time
+import warnings
 from typing import Any, Dict, List, Tuple
 from .ci_helpers import (
     check_for_discrepancies_and_log_everything_into_a_json_file,
@@ -97,7 +98,6 @@ def get_untrained_model(model_id: str, second_input: bool, verbose: int) -> Dict
                 },
                 # "_attn_implementation": "flash_attention_2",
                 "_attn_implementation": "sdpa",
-                "dtype": "float16",
             }
 
     config_reduction = _config_reduction
@@ -281,6 +281,10 @@ def main(
         ).eval()
         data = dict(model=model)
         config = model.config
+        if not hasattr(config, "bos_token_id") or not config.bos_token_id:
+            config.bos_token_id = 151643
+        if not hasattr(config, "eos_token_id") or not config.eos_token_id:
+            config.eos_token_id = 151645
     else:
         print("-- random model")
         data = get_untrained_model(model_id, second_input=second_input, verbose=1)
@@ -298,7 +302,11 @@ def main(
     print(f"-- config._attn_implementation={model.config._attn_implementation}")
     print(f"-- model.dtype={model.dtype}")
     print(f"-- model.device={model.device}")
-    processor = AutoProcessor.from_pretrained(model_id, use_fast=True)
+    try:
+        processor = AutoProcessor.from_pretrained(model_id, use_fast=True)
+    except OSError as e:
+        warnings.warn(f"Unable to access internet due to {e!r}", ResourceWarning, stacklevel=0)
+        return
     print(f"-- processor={type(processor)}")
 
     export_inputs, other_inputs = None, None

onnx_diagnostic/export/api.py

@@ -1,6 +1,11 @@
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+import inspect
+import os
+import textwrap
+from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, Union
 import torch
+from .dynamic_shapes import ModelInputs
 from .onnx_plug import EagerDirectReplacementWithOnnx
+from ..helpers import string_type
 
 
 def get_main_dispatcher(
@@ -70,6 +75,7 @@ def to_onnx(
     inline: bool = True,
 ) -> Any:
     """
+    Exports one model into ONNX.
     Common API for exporters. By default, the models are optimized to use the
     most efficient kernels implemented in :epkg:`onnxruntime`.
 
@@ -126,8 +132,12 @@ def to_onnx(
         from experimental_experiment.xbuilder import OptimizationOptions
 
         options = None
+        export_options = None
         if exporter_kwargs is not None:
             options = exporter_kwargs.pop("options", None)
+            export_options = exporter_kwargs.pop("export_options", None)
+        if export_options is None:
+            export_options = ExportOptions(save_ep=save_ep)
         if options is None and optimize:
             options = OptimizationOptions(
                 patterns="default+onnxruntime" if optimizer_for_ort else "default"
@@ -138,7 +148,7 @@ def to_onnx(
             else None
         )
 
-        return _to_onnx(
+        proto, opt_stats = _to_onnx(
             mod,
             args=args,
             kwargs=kwargs,
@@ -150,15 +160,52 @@ def to_onnx(
             dynamic_shapes=dynamic_shapes,
             large_model=True,
             output_dynamic_shapes=output_dynamic_shapes,
-            export_options=ExportOptions(save_ep=save_ep),
+            export_options=export_options,
             options=options,
             inline=inline,
             dispatcher=main_dispatcher,
+            optimize=optimize,
+            return_optimize_report=True,
             **(exporter_kwargs or {}),
         )
+        if opt_stats and filename and os.path.exists(filename):
+            import pandas
+
+            stat_filename = f"{os.path.splitext(filename)[0]}.opt.xlsx"
+            pattern_stats = []
+            for k, v in opt_stats.items():
+                if "time" in k:
+                    pattern_stats.append(dict(level="main", pattern=k, time_in=v))
+            pattern_stats.extend(
+                [{**obs, "level": "detailed"} for obs in opt_stats["optimization"]]
+            )
+            df = pandas.DataFrame(pattern_stats)
+            df.to_excel(stat_filename, index=False)
+            cols = [
+                c
+                for c in [
+                    "level",
+                    "pattern",
+                    "time_in",
+                    "iteration",
+                    "inlined",
+                    "removed",
+                    "added",
+                    "instances",
+                    "changed",
+                    "scale",
+                ]
+                if c in df.columns
+            ]
+            agg = {k: "sum" for k in cols if k not in ("level", "pattern")}
+            agg.update(dict(iteration="max", instances="mean"))
+            agg = {k: v for k, v in agg.items() if k in df.columns}
+            stat_filename = f"{os.path.splitext(filename)[0]}.opt.agg.xlsx"
+            df[cols].groupby(["level", "pattern"]).agg(agg).to_excel(stat_filename)
+
+        return proto
 
     if exporter in ("dynamo", "onnx-dynamo"):
-        import os
         from ..helpers import flatten_object
         import onnxscript.rewriter.ort_fusions as ort_fusions
 
@@ -225,7 +272,6 @@ def to_onnx(
         return epo
 
     if exporter == "modelbuilder":
-        import os
         from ..helpers import flatten_object, string_type
         from ..helpers.model_builder_helper import create_model_builder, save_model_builder
 
@@ -266,3 +312,247 @@ def to_onnx(
         return onx
 
     raise ValueError(f"Unknown exporter={exporter!r}")
+
+
+class _WrapperToExportMethodToOnnx(torch.nn.Module):
+    """
+    Wraps an existing models in order to spy on inputs.
+    This is used by :func:`onnx_diagnostic.export.api.method_to_onnx`.
+    """
+
+    def __init__(
+        self,
+        mod: "torch.nn.Module",
+        method_name: str = "forward",
+        input_names: Optional[Sequence[str]] = None,
+        target_opset: Optional[Union[int, Dict[str, int]]] = None,
+        verbose: int = 0,
+        filename: Optional[str] = None,
+        output_names: Optional[List[str]] = None,
+        output_dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
+        exporter: str = "onnx-dynamo",
+        exporter_kwargs: Optional[Dict[str, Any]] = None,
+        save_ep: Optional[str] = None,
+        optimize: bool = True,
+        optimizer_for_ort: bool = True,
+        use_control_flow_dispatcher: bool = False,
+        onnx_plugs: Optional[List[EagerDirectReplacementWithOnnx]] = None,
+        inline: bool = True,
+        convert_after_n_calls: int = 2,
+        patch_kwargs: Optional[Dict[str, Any]] = None,
+        skip_kwargs_names: Optional[Set[str]] = None,
+        dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
+    ):
+        super().__init__()
+        self._model_to_call = mod
+        self._method_name = method_name
+        self._method_call = (
+            self._model_to_call.forward
+            if method_name == "forward"
+            else getattr(mod, method_name)
+        )
+        self._inputs: List[Tuple[Tuple[Any, ...], Dict[str, Any]]] = []
+        self._convert_after_n_calls = convert_after_n_calls
+        self._patch_kwargs = patch_kwargs
+        self._method_src = None
+        self.verbose = verbose
+        self.skip_kwargs_names = skip_kwargs_names
+        self.dynamic_shapes = dynamic_shapes
+        self._to_onnx_kwargs = dict(
+            input_names=input_names,
+            target_opset=target_opset,
+            verbose=verbose,
+            filename=filename,
+            output_names=output_names,
+            output_dynamic_shapes=output_dynamic_shapes,
+            exporter=exporter,
+            exporter_kwargs=exporter_kwargs,
+            save_ep=save_ep,
+            optimize=optimize,
+            optimizer_for_ort=optimizer_for_ort,
+            use_control_flow_dispatcher=use_control_flow_dispatcher,
+            onnx_plugs=onnx_plugs,
+            inline=inline,
+        )
+        self._export_done = False
+
+    def __str__(self) -> str:
+        return self.__repr__()
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}({self._model_to_call.__class__.__name__}."
+            f"{self._method_name})"
+        )
+
+    def forward(self, *args, **kwargs):
+        if not self._export_done:
+            self._inputs.append(
+                (
+                    args,
+                    (
+                        kwargs
+                        if not kwargs or not self.skip_kwargs_names
+                        else {
+                            k: v for k, v in kwargs.items() if k not in self.skip_kwargs_names
+                        }
+                    ),
+                )
+            )
+            if self.verbose:
+                print(
+                    f"[method_to_onnx] input[{len(self._inputs)-1}]: "
+                    f"{string_type(self._inputs[-1], with_shape=True)}"
+                )
+            if len(self._inputs) >= self._convert_after_n_calls:
+                self._convert_method_to_onnx()
+                del self._inputs[:]
+                self._export_done = True
+        return self._method_call(*args, **kwargs)
+
+    def _convert_method_to_onnx(self):
+
+        def make_method(self):
+            inner_sig = inspect.signature(self._method_call)
+            params = [
+                p.replace(annotation=inspect._empty) for p in inner_sig.parameters.values()
+            ]
+            simple_sig = inspect.Signature(params, return_annotation=inspect._empty)
+            args = str(simple_sig)[1:-1]
+            calls_args = ", ".join(f"{p}={p}" for p in simple_sig.parameters)
+            src = textwrap.dedent(
+                f"""
+                def f(self, {args}):
+                    return self._method_call({calls_args})
+                """
+            )
+            self._method_src = src
+            ns = {}
+            try:
+                exec(src, ns)
+            except NameError as e:
+                raise NameError(f"Unable to compile due to {e}\n{src}") from e
+            return ns["f"]
+
+        class WrapWithExactSignature(torch.nn.Module):
+            def __init__(self, parent):
+                super().__init__()
+                self._model_to_call = parent._model_to_call
+                self._method_call = parent._method_call
+
+            forward = make_method(self)
+
+        compiled_model = WrapWithExactSignature(self)
+
+        if self.dynamic_shapes is None:
+            mi = ModelInputs(compiled_model, self._inputs)
+            ds = mi.guess_dynamic_shapes()
+            if self.verbose:
+                print(f"[method_to_onnx] guess_dynamic_shapes={string_type(ds)}")
+            a, kw, nds = mi.move_to_kwargs(*self._inputs[-1], ds)
+        else:
+            a, kw = self._inputs[-1]
+            nds = [self.dynamic_shapes]
+        if self.verbose:
+            print(f"[method_to_onnx] export args={string_type(a, with_shape=True)}")
+            print(f"[method_to_onnx] export kwargs={string_type(kw, with_shape=True)}")
+            print(f"[method_to_onnx] dynamic_shapes={string_type(nds)}")
+        if self._patch_kwargs is None:
+            to_onnx(
+                compiled_model,
+                args=a,
+                kwargs=kw,
+                dynamic_shapes=nds[-1],
+                **self._to_onnx_kwargs,
+            )
+            return
+        from ..torch_export_patches import torch_export_patches
+
+        with torch_export_patches(**self._patch_kwargs):
+            to_onnx(
+                compiled_model,
+                args=a,
+                kwargs=kw,
+                dynamic_shapes=nds[-1],
+                **self._to_onnx_kwargs,
+            )
+
+
+def method_to_onnx(
+    mod: "torch.nn.Module",
+    method_name: str = "forward",
+    input_names: Optional[Sequence[str]] = None,
+    target_opset: Optional[Union[int, Dict[str, int]]] = None,
+    verbose: int = 0,
+    filename: Optional[str] = None,
+    output_names: Optional[List[str]] = None,
+    output_dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
+    exporter: str = "onnx-dynamo",
+    exporter_kwargs: Optional[Dict[str, Any]] = None,
+    save_ep: Optional[str] = None,
+    optimize: bool = True,
+    optimizer_for_ort: bool = True,
+    use_control_flow_dispatcher: bool = False,
+    onnx_plugs: Optional[List[EagerDirectReplacementWithOnnx]] = None,
+    inline: bool = True,
+    convert_after_n_calls: int = 2,
+    patch_kwargs: Optional[Dict[str, Any]] = None,
+    skip_kwargs_names: Optional[Set[str]] = None,
+    dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None,
+) -> Callable:
+    """
+    Exports one method into ONNX for a module into ONNX.
+    It returns a new method which must be called by the user
+    at least twice with different values for the dynamic dimension
+    between triggering the conversion into ONNX.
+
+    :param mod_meth: function to export into ONNX
+    :param input_names: input names for the onnx model (optional)
+    :param target_opset: opset to target, if not specified, each converter
+        keeps its default value
+    :param verbose: verbosity level
+    :param filename: output filename, mandatory, the onnx model is saved on disk
+    :param output_names: to change the output of the onnx model
+    :param output_dynamic_shapes: to overwrite the dynamic shapes names
+    :param exporter: exporter to use (``onnx-dynamo``, ``modelbuilder``, ``custom``)
+    :param exporter_kwargs: additional parameters sent to the exporter
+    :param save_ep: saves the exported program
+    :param optimize: optimizes the model
+    :param optimizer_for_ort: optimizes the model for onnxruntime
+    :param use_control_flow_dispatcher: use the dispatcher created to supported
+        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
+    :param convert_after_n_calls: converts the model after this number of calls.
+    :param patch_kwargs: patch arguments
+    :param skip_kwargs_names: use default values for these parameters part of
+        the signature of the method to export
+    :param dynamic_shapes: dynamic shapes to use if the guessed ones are not right
+    :return: the output of the selected exporter, usually a structure including
+        an onnx model
+
+    See :ref:`l-plot-tiny-llm-export-method-generate` for an example.
+    """
+    wrapped_model = _WrapperToExportMethodToOnnx(
+        mod=mod,
+        method_name=method_name,
+        input_names=input_names,
+        target_opset=target_opset,
+        verbose=verbose,
+        filename=filename,
+        output_names=output_names,
+        output_dynamic_shapes=output_dynamic_shapes,
+        exporter=exporter,
+        exporter_kwargs=exporter_kwargs,
+        save_ep=save_ep,
+        optimize=optimize,
+        optimizer_for_ort=optimizer_for_ort,
+        use_control_flow_dispatcher=use_control_flow_dispatcher,
+        onnx_plugs=onnx_plugs,
+        inline=inline,
+        convert_after_n_calls=convert_after_n_calls,
+        patch_kwargs=patch_kwargs,
+        skip_kwargs_names=skip_kwargs_names,
+        dynamic_shapes=dynamic_shapes,
+    )
+    return wrapped_model

onnx_diagnostic/export/cf_simple_loop_for.py

@@ -11,6 +11,7 @@ from torch._higher_order_ops.utils import (
     unique_graph_id,
     validate_subgraph_args_types,
 )
+import torch._dynamo.variables.higher_order_ops as hop
 from torch.fx.experimental.proxy_tensor import ProxyTorchDispatchMode, track_tensor_tree
 from torch.utils._python_dispatch import _get_current_dispatch_mode
 
@@ -97,14 +98,18 @@ def _simple_loop_for_fn(
                 f"Unexpected number of results {len(r)} for function {body_fn}, "
                 f"expected {len(res[-1])}"
             )
+            assert all(isinstance(t, torch.Tensor) for t in r), (
+                f"Unexpected type {[type(_) for _ in r]} for returned by function {body_fn}, "
+                f"it must be a tuple of Tensor or a Tensor."
+            )
             res.append(r)
         else:
             assert isinstance(r, torch.Tensor), (
-                f"Unexpected type {r} for function {body_fn}, "
-                f"it must be a tuple or a Tensor."
+                f"Unexpected type {type(r)} coming from function {body_fn}, "
+                f"it must be a tuple of Tensor or a Tensor."
             )
             assert not res or len(res[-1]) == 1, (
-                f"Unexpected number of results {len(r)} for function {body_fn}, "
+                f"Unexpected number of results {len(r)} coming from function {body_fn}, "
                 f"expected {len(res[-1])}"
             )
             res.append((r,))
@@ -126,8 +131,6 @@ def _simple_loop_for_fn(
     )
 
 
-# from torch._functorch.utils import exposed_in
-# @exposed_in("torch")
 def _simple_loop_for(
     n_iter: Union[int, torch.Tensor],
     body_fn: Callable,
@@ -159,7 +162,7 @@ def _simple_loop_for(
 
     if torch.compiler.is_dynamo_compiling():
         return simple_loop_for_op(
-            n_iter, body_fn, (n_iter, *operands), concatenation_dims=concatenation_dims
+            n_iter, body_fn, operands, concatenation_dims=concatenation_dims
         )
 
     if isinstance(n_iter, (bool, int, float)):
@@ -181,8 +184,10 @@ def _simple_loop_for(
 
     with setup_compilation_env() as _backend:
         return _loop_for_op_wrapper(n_iter, body_fn, operands, concatenation_dims)
-        # return torch.compile(_loop_for_op_wrapper, backend=backend, fullgraph=True)(
-        #    n_iter, body_fn, operands, concatenation_dims)
+        # This is needed to support function body using module weights or function body
+        # defined as a class method. This is yet to be implemented.
+        # cpl = torch.compile(_loop_for_op_wrapper, backend=_backend, fullgraph=True)
+        # return cpl(n_iter, body_fn, operands, concatenation_dims)
 
 
 def trace_simple_loop_for(
@@ -236,9 +241,15 @@ def loop_for_op_dense(n_iter, body_fn, operands, concatenation_dims=None):
     )
     mode = _get_current_dispatch_mode()
     assert mode is None, "Mode should never be enabled for CPU/CUDA key"
-    return _simple_loop_for_fn(
-        n_iter, body_fn, operands, concatenation_dims=concatenation_dims
+    is_fake = isinstance(n_iter, torch._subclasses.fake_tensor.FakeTensor)
+    res = _simple_loop_for_fn(n_iter, body_fn, operands, concatenation_dims=concatenation_dims)
+    assert is_fake or not any(
+        isinstance(r, torch._subclasses.fake_tensor.FakeTensor) for r in res
+    ), (
+        f"One result is a fake tensor but the inputs were not, type(n_iter)={type(n_iter)}, "
+        f"operands: {[type(_) for _ in operands]}, res: {[type(_) for _ in res]}"
     )
+    return res
 
 
 @simple_loop_for_op.py_impl(ProxyTorchDispatchMode)
@@ -267,6 +278,180 @@ simple_loop_for_op.fallthrough(torch._C.DispatchKey.AutogradCPU)
 simple_loop_for_op.fallthrough(torch._C.DispatchKey.AutogradCUDA)
 
 
+class SimpleLoopForHigherOrderVariable(hop.TorchHigherOrderOperatorVariable):
+    """
+    Replicates the same pattern found for other higher order operators.
+    This enables recursive compilation and the use of modules inside a function.
+    """
+
+    _HOP_NAME = "simple_loop_for"
+    _ALLOW_FALLBACK_TO_EAGER = False
+    supports_input_mutation = False
+    supports_aliasing = False
+
+    def _call_function(
+        self,
+        tx: torch._dynamo.symbolic_convert.InstructionTranslator,
+        args: list[hop.VariableTracker],
+        kwargs: dict[str, hop.VariableTracker],
+    ) -> hop.VariableTracker:
+        """Main function."""
+        args, kwargs = hop.LazyVariableTracker.realize_all((args, kwargs))
+
+        for i, k in enumerate(["n_iter", "body_fn", "operands", "concatenated_dims"]):
+            if v := kwargs.pop(k, None):
+                assert i == len(args), "did not provide the right number of non-keyword args"
+                args.append(v)
+
+        if len(args) != 4 or kwargs:
+            hop.unimplemented(
+                gb_type="simple_loop_for: improper args/kwargs",
+                context=f"args: {args}, kwargs: {kwargs}",
+                explanation=f"torch.cond expects 4 positional arguments (got {len(args)}) "
+                f"and no keyword arguments (got {len(kwargs)})",
+                hints=[*hop.graph_break_hints.USER_ERROR],
+            )
+
+        # Specialize into one of the branches since pred is constant
+        n_iter, body_fn, operands, _concatenated_dims = args
+        assert type(n_iter) is not hop.ConstantVariable, (
+            f"n_iter is a {type(n_iter)}. When used simple_loop_for, "
+            f"it unrolls the loop. A SymInt should be used."
+        )
+
+        # predicate
+        if type(n_iter.realize()) not in (
+            hop.ConstantVariable,
+            hop.TensorVariable,
+            hop.SymNodeVariable,
+        ):
+            hop.unimplemented(
+                gb_type="simple_loop_for: improper predicate",
+                context=str(n_iter),
+                explanation=(
+                    f"Expected `n_iter` to be an int or a integer "
+                    f"tensor with a single item "
+                    f"but got {str(type(n_iter))} with original python type "
+                    f"{str(n_iter.python_type())}."
+                ),
+                hints=[*hop.graph_break_hints.USER_ERROR],
+            )
+
+        # operands
+        if not isinstance(operands, (hop.ListVariable, hop.TupleVariable)):
+            hop.unimplemented(
+                gb_type="simple_loop_for: improper operands",
+                context=str(operands),
+                explanation="Expected `operands` to be a list/tuple "
+                f"but got {operands.python_type()}.",
+                hints=[*hop.graph_break_hints.USER_ERROR],
+            )
+
+        operands_seq = operands.unpack_var_sequence(tx)
+        if not hop.only_consist_of(
+            operands, (hop.TensorVariable, hop.ConstantVariable, hop.SymNodeVariable)
+        ):
+            hop.unimplemented(
+                gb_type="simple_loop_for: improper operands contents",
+                context=str(operands),
+                explanation=(
+                    "Expected `operands` to be a list/tuple of pytrees "
+                    "that only consists of tensor leaves."
+                ),
+                hints=[*hop.graph_break_hints.USER_ERROR],
+            )
+
+        # branches
+        hop._check_supported_callable_arg(tx, body_fn, "body_fn")
+
+        def speculate_body():
+            (
+                (ret_val, ret_spec),
+                ret_graph,
+                ret_lifted_freevars,
+            ) = hop.speculate_subgraph(
+                tx,
+                args[1],
+                (args[0], *operands_seq),
+                {},
+                self._HOP_NAME,
+                source_target=self.value,
+                should_flatten_outputs=True,
+                # TODO - removing consts from control flow ops need more work
+                remove_consts_from_outputs=False,
+                supports_input_mutation=self.supports_input_mutation,
+                supports_aliasing=self.supports_aliasing,
+            )
+
+            # need to ensure we increase epoch so we don't memoize unbacked bindings
+            # across different subgraphs which can interfere with runtime assertion
+            # generation.
+            tx.fake_mode.epoch += 1
+
+            if not hop.only_consist_of(ret_val, (hop.TensorVariable, hop.ConstantVariable)):
+                hop.unimplemented(
+                    gb_type="simple_loop_for: unsupported branch return type",
+                    context=str(ret_val),
+                    explanation=(
+                        "Expected branches to return a possibly nested "
+                        "pytree of tensors or constant ints."
+                    ),
+                    hints=[*hop.graph_break_hints.USER_ERROR],
+                )
+            for ret in ret_val.unpack_var_sequence(tx):
+                if ret.is_python_constant() and not isinstance(ret.as_python_constant(), int):
+                    hop.unimplemented(
+                        gb_type=(
+                            "simple_loop_for: unsupported branch return type "
+                            "(constant non-int)"
+                        ),
+                        context=str(ret_val),
+                        explanation="Constants returned from branches must be ints.",
+                        hints=[*hop.graph_break_hints.USER_ERROR],
+                    )
+            return ret_val, ret_spec, ret_graph, ret_lifted_freevars
+
+        body_r, body_spec, body_graph, body_lifted_freevars = speculate_body()
+        body_nn_modules = dict(tx.output.nn_modules)
+
+        same_spec = body_spec.treespec.as_python_constant()
+        if same_spec is not NotImplemented and not same_spec:
+            hop.unimplemented(
+                gb_type="simple_loop_for: differing branch outputs",
+                context=(
+                    f"body_spec: {body_spec.treespec}, false_spec: "
+                    f"{body_spec.treespec}, same_spec: {same_spec}"
+                ),
+                explanation="Expected branches to return the same pytree structure.",
+                hints=[*hop.graph_break_hints.USER_ERROR],
+            )
+
+        body_name = tx.output.install_subgraph(
+            "loop_body", torch.fx.GraphModule(body_nn_modules, body_graph)
+        )
+        body_node = hop.make_attr(tx, body_name)
+        p_args = (
+            n_iter.as_proxy(),
+            body_node,
+            # We pick true_shared but it shouldn't matter
+            operands.as_proxy() + tuple(body_lifted_freevars.keys()),
+        )
+
+        return hop._call_function_and_unflatten_output(
+            tx,
+            simple_loop_for,
+            p_args,
+            {},
+            None,
+            body_spec,
+            body_r,
+        )
+
+
+hop._hop_name_to_variable_class["simple_loop_for"] = SimpleLoopForHigherOrderVariable
+
+
+# @torch._functorch.utils.exposed_in("torch")
 def simple_loop_for(
     n_iter: Union[int, torch.Tensor],
     body_fn: Callable,