(while_loop)=

# Control Flow - While Loop

`torch.while_loop` is a structured control flow operator that runs a body function while a condition is true.
It can logically be seen as implemented as follows:

```python
def while_loop(
    cond_fn: Callable[..., bool],
    body_fn: Callable[..., tuple],
    carried_inputs: tuple,
):
    val = carried_inputs
    while cond_fn(*val):
        val = body_fn(*val)
    return val
```

```{warning}
`torch.while_loop` is a prototype feature in PyTorch. It has limited support for input and output types.
Please look forward to a more stable implementation in a future version of PyTorch.
Read more about feature classification at: https://pytorch.org/blog/pytorch-feature-classification-changes/#prototype
```

## Examples

Below is a basic example that uses while_loop to iterate until a condition is met:

```{code-cell}
import torch
from torch._higher_order_ops import while_loop

class M(torch.nn.Module):

    def cond_fn(self, iter_count, x):
        return iter_count.sum() > 0

    def body_fn(self, iter_count, x):
        return iter_count - 1, x * 2

    def forward(self, init_iter, init_x):
        final_iter, final_x = while_loop(self.cond_fn, self.body_fn, (init_iter, init_x))
        return final_iter, final_x

m = M()
```

We can eagerly run the model and expect the results vary based on input shape:

```{code-cell}
_, final_x = m(torch.tensor([3]), torch.ones(3))
assert torch.equal(final_x, torch.ones(3) * 2**3)

_, final_x = m(torch.tensor([10]), torch.ones(3))
assert torch.equal(final_x, torch.ones(3) * 2**10)
```

We can export the model for further transformations and deployment. This gives us an exported program that preserves the while_loop structure:

```{code-cell}
ep = torch.export.export(M(), (torch.tensor([10]), torch.ones(3)))
print(ep)
```

Notice that both the condition and body functions become sub-graph attributes of
the top-level graph module.

## Restrictions

- `body_fn` must return tensors or integers with the same metadata (shape, dtype) as inputs.

- `body_fn` and `cond_fn` must not in-place mutate the `carried_inputs`. A clone before mutation is required.

- `body_fn` and `cond_fn` must not mutate Python variables (e.g., list/dict) created outside the function.

- `body_fn` and `cond_fn`'s output cannot alias any of the inputs. A clone is required.


## API Reference

```{eval-rst}
.. autofunction:: torch._higher_order_ops.while_loop.while_loop
```