Repository: pytorch
Issue: Tracing models all the way down to basic functions
Labels: feature,triaged,module: fx

## 🚀 Feature

<!-- A clear and concise description of the feature proposal -->

I propose the enhancement of `torch.fx.Tracer` such that it is able to trace models all the way down to basic torch functions (like the ones in `torch.functional` or `torch.nn.functional`, as opposed to stopping at predefined modules in `torch.nn`).

## Motivation

<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->

In an attempt to extract the flat computational graph consisting of only basic function calls (not containing calls to function-like stateful objects such as `torch.nn.Module`) from a PyTorch model, I tried to use `torch.fx.Tracer`. It enables customization of what constitutes a leaf module through its `is_leaf_module` method. Overriding that method returning `False` for all modules would result in tracing models all the way to basic function calls, which is clearly one of the most basic use cases of a tracer.

However, the current `Tracer` breaks down even on simple models, such as `torchvision.models.resnet18`. The problem in this case comes from `torch.nn.BatchNorm2d` module (but could be similarly present in other modules as well), which contains a check on the input's shape (`if input.dim() != 4:`), which the tracer cannot evaluate, since it does not have shape information about tensors.

I have also tried using the `concrete_args` parameter of the `trace` method, but as far as I understand, it would result in treating tensor inputs as constants and folding the computation, so the resulting trace would be meaningless (just a constant node), even if the trace would not fail (however, it does even in this case, which I am not sure if it's a bug or not).

(Note that the function `torch.fx.symbolic_trace` does work on the above model (as shown in one of the examples), as it does not try to trace all the way down to basic functions (stops at module level), and hence does not hit the above mentioned check, but that is just avoiding the problem, not solving it.)

Enabling such tracing would be essential for being able to extract and convert torch models for execution in other inference engines (for example running on custom hardware).

## Pitch

<!-- A clear and concise description of what you want to happen. -->

I'd like to have a way of tracing models all the way down to basic functions, given that the model does not contain true control flow that depends on actual input _data_.

Ideally, in general, I'd like the tracer to be able to trace models given some concrete inputs that give the tracer information about tensor shapes. All control flow / checks that depend only on tensor shapes that can be derived from example inputs should be evaluated during tracing, and the traced model should be valid for inputs whose shape is identical to those of the example inputs given to the tracer.

However, for the exact problem with the above model, even less would be sufficient. As I see, such sanity checks on tensor shapes play no role in the actual computation, but at the same time they are blocking the tracer from tracing the actual computation graph. Resolving this somehow could also be a first step, for example such checks could be annotated somehow such that the tracer could omit them (as they play no computational role in the trace) or simply copy them to the traced graph.

## Alternatives

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->

I have also considered using `torch.jit.trace_module`. Since that function receives example inupts, it can derive shape information from that, and so it can trace through such control flow / checks. However, the traced graph's details cannot be publicly accessed, it cannot be manipulated (with graph transformations, for example constant folding) as in `torch.fx`. In particular, shape information does not seem to be stored in the graph, and shape propagation cannot be performed as with `torch.fx`.

## Additional context

<!-- Add any other context or screenshots about the feature request here. -->

The problem can be reproduced with the following code:

```
import torch
import torch.fx as fx
import torchvision.models as models

class LeafTracer(fx.Tracer):

    def is_leaf_module(self, module, name) -> bool:
        return False

model = models.resnet18()
traced = LeafTracer().trace(model)
```

It fails with

```
torch.fx.proxy.TraceError: symbolically traced variables cannot be used as inputs to control flow
```

cc @ezyang @mruberry @jbschlosser