Export and On-Device Runtime

## What export is trying to do `torch.export.export` produces one static graph that describes the model for every input it will ever see. It does that by running your `forward()` once with example inputs while substituting symbolic tensors for real data. Any Python construct that needs to *look at the data* to decide what to do next is therefore a problem: at trace time the data is not there, and at run time the Python is not there. This is the sharpest difference between export and `torch.compile`. `torch.compile` may hit something it cannot capture, insert a **graph break**, and let plain Python execute that region before resuming compiled code. Export has no such escape hatch, and neither does the ExecuTorch runtime on the device — it walks a serialized instruction stream. Full-graph capture is mandatory. ## The two failure modes Write this in a module: ``` def forward(self, x): if x.sum() > 0: return self.a(x) return self.b(x) ``` `x.sum() > 0` is a zero-dimensional tensor. Python's `if` calls `__bool__` on it, which forces a concrete answer. What happens next depends on how the value was produced: 1. **Loud failure.** If the condition depends on a value the exporter is tracking symbolically, export raises a data-dependent error — it cannot guard on an expression whose truth value is unknown. This is the good outcome: you find out on your build machine. 2. **Silent specialization.** If the condition collapses to a plain Python `bool` during tracing — for instance because it depends on a shape that got specialized, on a configuration attribute, or on a constant — export happily records only the taken branch. The other branch simply does not exist in the graph. The exported model runs, is fast, and is wrong for half your inputs. The second is the classic on-device bug, because the eager model in your notebook and the `.pte` on the phone disagree only for inputs your example did not cover. It is also why output-parity testing across a *varied* input set, not one sample, belongs in the export CI job. ## The same trap in loops `for i in range(x.shape[0])` behaves the same way. If the batch dimension is specialized, the loop is unrolled to exactly that many iterations and baked into the graph; feed a different batch size and the guard rejects it (or, worse, the model quietly processes the wrong number of elements if you built the loop over a Python integer you also passed in). `while` loops with tensor conditions fail outright. ## Fix 1 — `torch.cond` `torch.cond(pred, true_fn, false_fn, operands)` is the supported way to keep a genuine runtime branch: ``` return torch.cond(x.sum() > 0, self.a, self.b, (x,)) ``` It is a higher-order operator: both branches are traced and both end up in the graph, and the predicate is evaluated at run time. The constraints are real and worth stating in an interview: both branches must take the same operands, return the same number of outputs with matching dtypes and compatible shapes, must not mutate their inputs or close over tensors outside `operands`, and each must itself be exportable. Backend delegation is also less complete inside conditionals than in straight-line code, so a `torch.cond` region may execute on the runtime's own kernels rather than the accelerator. ## Fix 2 — make it arithmetic Many "branches" are really elementwise selection. `torch.where(cond, a, b)` produces a data-dependent *result* with no control flow at all. The cost is that both sides are always computed, so this is right for cheap expressions and wrong for skipping an expensive subnetwork. ## Fix 3 — hoist the decision out of the model Often the cleanest answer, and the one senior candidates reach for: if the branch decides between two genuinely different networks, export **two** methods (or two `.pte` files) and let the application code pick. The branch condition becomes app logic in Kotlin or Swift, where branching is free and debuggable, and each exported graph stays straight-line and fully delegatable. Preprocessing decisions — orientation, language, model size — usually belong here rather than inside `forward()`. ## Practical shape of the answer When an interviewer hands you "this model won't export", the useful sequence is: read the error to see whether it is a data-dependent guard or an unsupported operator; find the offending `if`/`while` in `forward()`; decide whether the branch is *semantic* (needs `torch.cond` or two exported methods) or *incidental* (a debug path, a training-only branch, an assertion) and simply delete it for the export path; then re-export and diff outputs against eager on inputs that exercise **both** sides of the original branch.