pytorch / pytorch/pytorch

[feature request] Re-enable `__torch_function__` dispatch when invoking autograd hooks for Tensor subclasses

Open
#167,223 7 comments 0 reactions 0 assignees View on GitHub
actionable module: __torch_function__ module: autograd module: docs triaged
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.