add autowrap_functions kwarg to fx.Tracer (#62106)
authorAlexander Soare <alexander.soare159@gmail.com>
Fri, 13 Aug 2021 00:35:02 +0000 (17:35 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Fri, 13 Aug 2021 00:38:25 +0000 (17:38 -0700)
commit219ba6575b682a9b61476da041c2220142d20e3b
tree596eaa67224434e97d66a87ccbaa1bb062969e6f
parent7a1ab9f5d730d2d329c59efa0dcc06ab5ed84fe7
add autowrap_functions kwarg to fx.Tracer (#62106)

Summary:
Implements feature request https://github.com/pytorch/pytorch/issues/62021

Test it out with

```python
from torch import fx
from torch import nn

def fx_int(x):
    return int(x)

class MyModule(nn.Module):
    def forward(self, x):
        return fx_int(x.shape[0] / 2)

tracer = fx.Tracer(autowrap_functions=(fx_int,))  # or remove kwarg to demonstrate symbolic trace error
tracer.trace(MyModule())
```

First time contributor, so please advise if I could have done anything to make lives easier for next time.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/62106

Reviewed By: SplitInfinity, driazati

Differential Revision: D30080834

Pulled By: jamesr66a

fbshipit-source-id: 68fadf8c881ea7930e7afd62b642874010fe4903
test/test_fx.py
torch/fx/_symbolic_trace.py