Add __torch_function__ for methods

[hameerabbasi]

According to pytorch/rfcs#3

From the goals in the RFC:

1. Support subclassing torch.Tensor in Python (done here)
2. Preserve torch.Tensor subclasses when calling torch functions on them (done here)
3. Use the PyTorch API with torch.Tensor-like objects that are not torch.Tensor
   subclasses (done in Add the __torch_function__ API override mechanism #30730)
4. Preserve torch.Tensor subclasses when calling torch.Tensor methods. (done here)
5. Propagating subclass instances correctly also with operators, using
   views/slices/indexing/etc. (done here)
6. Preserve subclass attributes when using methods or views/slices/indexing. (done here)
7. A way to insert code that operates on both functions and methods uniformly
   (so we can write a single function that overrides all operators). (done here)
8. The ability to give external libraries a way to also define
   functions/methods that follow the __torch_function__ protocol. (will be addressed in a separate PR)

This PR makes the following changes:

1. Adds the self argument to the arg parser.
2. Dispatches on self as well if self is not nullptr.
3. Adds a torch._C.DisableTorchFunction context manager to disable __torch_function__.
4. Adds a torch::torch_function_enabled() and torch._C._torch_function_enabled() to check the state of __torch_function__.
5. Dispatches all torch._C.TensorBase and torch.Tensor methods via __torch_function__.

TODO:

• Sequence Methods
• Docs
• Tests

Closes #28361

Benchmarks in #37091 (comment)

[dr-ci]

💊 CI failures summary and remediations

As of commit e715d90 (more details on the Dr. CI page):

---

💚 💚 Looks good so far! There are no failures yet. 💚 💚

---

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions on the GitHub issue tracker or post in the (internal) Dr. CI Users group.

See how this bot performed.

This comment has been revised 442 times.

[ezyang]

Urrgh, there was a force push to the branch, this is going to make keeping the internal FB version in sync substantially more complicated (as I cannot just cherry pick the latest changes from open source)

[hameerabbasi]

Urrgh, there was a force push to the branch, this is going to make keeping the internal FB version in sync substantially more complicated (as I cannot just cherry pick the latest changes from open source)

I apologize, I only force-pushed the last commit: 25f37f3

[ezyang]

@ezyang should we update the torch_function for https://pytorch.org/blog/pytorch-feature-classification-changes/? What would you say is the classification, Beta

I need to submit it to the internal classification process. But I would agree it's somewhere between prototype and beta; I think __torch_function__ as the original behavior clearly hits the bar for beta, but the new method functionality is much more prototype-y.

[ezyang]

Aug 03 11:59:59 + python -c 'import torch; print(torch.__config__.show())'
Aug 03 11:59:59 Traceback (most recent call last):
Aug 03 11:59:59   File "<string>", line 1, in <module>
Aug 03 11:59:59   File "/opt/conda/lib/python3.8/site-packages/torch/__init__.py", line 526
Aug 03 11:59:59     quantized_gru = torch.ops.aten.quantized_gru§

test failures

[facebook-github-bot]

@ezyang has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[ezyang]

mypy failing now

[hameerabbasi]

Any further action needed in this PR?

[facebook-github-bot]

@ezyang has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[ezyang]

just tryin to land it 👍

[rgommers]

It landed! <does a little dance and crosses fingers>

Thanks @hameerabbasi and @ezyang!