[feature request] Re-enable `__torch_function__` dispatch when invoking autograd hooks for Tensor subclasses
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.5k
- PR merge metrics
- PR metrics pending
Description
### 🚀 The feature, motivation and pitch
### The feature:
When PyTorch invokes user-defined autograd hooks (e.g. `register_post_accumulate_grad_hook`), temporarily re-enable the `__torch_function__` dispatch for Tensor subclasses so that `__torch_function__` implementations on subclasses are called from inside hooks.
### Motivation / problem:
While experimenting with `__torch_function__` on both Tensor subclasses and Tensor-like objects, I observed that `__torch_function__` is not dispatched when calling functions in the torch namespace from inside autograd hooks for Tensor subclasses. For Tensor-like objects the dispatch works correctly inside hooks. This mismatch leads to surprising behaviour and forces subclass authors to manually wrap hook bodies in `with torch._C._EnableTorchFunction():` to get consistent dispatch.
### Current behaviour:
```py
class MyTensor(torch.Tensor):
def __new__(cls, data):
return torch.Tensor._make_subclass(cls, data, data.requires_grad)
@classmethod
def __torch_function__(cls, func, types, args=(), kwargs=None):
if kwargs is None:
kwargs = {}
print(f"[MyTensor __torch_function__] func={func.__name__}")
return super().__torch_function__(func, types, args, kwargs)
x = MyTensor(torch.randn(3, requires_grad=True))
y = MyTensor(torch.randn(3, requires_grad=True))
def post_accum_hook(param):
print("Start of hook")
param + MyTensor(torch.ones_like(param))
print("End of hook")
x.register_post_accumulate_grad_hook(post_accum_hook)
z = torch.add(x, y)
torch.sum(z).backward()
```
Output:
```
>>> [MyTensor __torch_function__] func=register_post_accumulate_grad_hook
>>> [MyTensor __torch_function__] func=add
>>> [MyTensor __torch_function__] func=sum
>>> [MyTensor __torch_function__] func=backward
>>> Start of hook
>>> End of hook
```
Compare with Tensor-like object (dispatch works inside hook):
```py
class MyTensorLike:
def __init__(self, data):
self.data = data
@classmethod
def __torch_function__(cls, func, types, args=(), kwargs=None):
if kwargs is None:
kwargs = {}
print(f"[MyTensorLike __torch_function__] func={func.__name__}")
result = func(*(arg.data if isinstance(arg, MyTensorLike) else arg for arg in args), **kwargs)
return cls(result)
a = MyTensorLike(torch.randn(3, requires_grad=True))
b = MyTensorLike(torch.randn(3, requires_grad=True))
def post_accum_grad_hook(param):
print("Start of hook")
param + MyTensorLike(torch.ones_like(param.data))
print("End of hook")
a.data.register_post_accumulate_grad_hook(post_accum_grad_hook)
c = torch.add(a, b)
torch.sum(c).data.backward()
```
Output:
```
>>> [MyTensorLike __torch_function__] func=add
>>> [MyTensorLike __torch_function__] func=sum
>>> Start of hook
>>> [MyTensorLike __torch_function__] func=add
>>> End of hook
```
Workaround: Manually enable dispatch in the hook:
```py
def post_accum_hook(param):
print("Start of hook")
with torch._C._EnableTorchFunction():
param + MyTensor(torch.ones_like(param))
print("End of hook")
```
This produces the expected `__torch_function__` print inside the hook.
### Root cause
When Tensor subclasses call `super().__torch_function__`, the `super()` implementation runs under [`with _C.DisableTorchFunctionSubclass():`](https://github.com/pytorch/pytorch/blob/7b423c2d217452d7f65788dc3a9cb786f0b45769/torch/_tensor.py#L1703) to prevent infinite recursion during dispatch. The autograd engine invokes hooks while that context is still in effect, so `__torch_function__` calls originating from inside the hook are suppressed for subclasses. Tensor-like objects do not call `super().__torch_function__`, so their dispatch remains enabled inside hooks.
### Alternatives
_No response_
### Additional context
_No response_
cc @svekars @sekyondaMeta @AlannaBurke @ezyang @albanD @gqchen @nikitaved @soulitzer @Varal7 @xmfan @hameerabbasi @rgommers
Contributor guide
Assessment
This issue has not been assessed yet.