TorchScript ignores user-defined subclasses of Tensor and custom implementations of `__torch_function__`

username: JakobHavtorn
cooked: <p>According to the <a href="https://pytorch.org/docs/stable/notes/extending.html" rel="noopener nofollow ugc">documentation</a>, PyTorch can be extended by creating a subclass of <code>torch.Tensor</code> and implementing the <code>__torch_function__</code> method. This works perfectly fine in standard “eager” execution mode, however, <strong>it does not support scripting</strong>. This is a major drawback since scripting, besides its inherent benefits, is also the preferred way to export models to ONNX.</p>
<p>The issue is twofold:</p>
<ol>
<li>Methods or attributes of the original <code>torch.Tensor</code> that are modified on a subclass of <code>torch.Tensor</code>, are silently ignored when used in scripted methods/modules. Similarly, any <em>new</em> methods or attributes of the <code>torch.Tensor</code> subclass (not originally on the <code>torch.Tensor</code>) result in a <code>RuntimeError</code> when calling <code>torch.jit.script</code> on a method or module calling them.</li>
<li>A Tensor subclass that has a custom <code>__torch_function__</code> implementation doesn’t get its <code>__torch_function__</code> scripted into the static graph. Instead, the <code>__torch_function__</code> is “outside” the graph and is called with the scripted function(s) as the <code>func</code> argument. This removes any custom overrides that follow the “Extending PyTorch” <a href="https://pytorch.org/docs/stable/notes/extending.html" rel="noopener nofollow ugc">documentation</a>.</li>
</ol>
<p>Issue 1 has previously been documented <a href="https://github.com/pytorch/pytorch/issues/51135#ref-issue-805652030" rel="noopener nofollow ugc">here</a>, but did not seem to gain much traction. I have not yet found any previous reports of issue 2.</p>
<p><strong>Illustration of issue 1</strong><br>
The script below is similar to the example from <a href="https://github.com/pytorch/pytorch/issues/51135#ref-issue-805652030" rel="noopener nofollow ugc">here</a>.</p>
<ul>
<li>
<p>The eager execution of <code>call_length()</code> prints 123456 as expected but the scripted version prints 2, matching the actual tensor but not the override.</p>
</li>
<li>
<p>When scripting <code>call_new_method()</code>, our <code>TensorSubClass</code> type is ignored and inferred to be a regular <code>torch.Tensor</code> which does not have a <code>new_method()</code> defined on it – and therefore scripting fails with a <code>RuntimeError</code>:</p>
<pre><code class="lang-auto">RuntimeError:
'Tensor' object has no attribute or method 'new_method'.:
  File "/home/jdh/repos/dreamstream/test_tensor_subclass.py", line 23
def call_new_method(x: TensorSubClass) -&gt; int:
    return x.new_method()
           ~~~~~~~~~~~~ &lt;--- HERE
</code></pre>
</li>
</ul>
<p><em>Issue 1 script</em></p>
<pre><code class="lang-python">import torch


class TensorSubClass(torch.Tensor):
    def __init__(self, tensor: torch.Tensor):
        super().__init__()
        self.tensor = tensor

    def __len__(self) -&gt; int:
        """Modify the behaviour of len() to return a constant value."""
        return 123456

    def new_method(self) -&gt; int:
        """Add a new method to the subclass that does not exist on `torch.Tensor`."""
        return 123456


def call_length(tensor: TensorSubClass) -&gt; int:
    return len(tensor)


def call_new_method(tensor: TensorSubClass) -&gt; int:
    return tensor.new_method()


torch.manual_seed(0)

call_length_scripted = torch.jit.script(call_length)

tensor = TensorSubClass(torch.rand((2, 3, 4)))

length_from_eager = call_length(tensor)
print("Eager: ", length_from_eager)  # Prints 123456
length_from_jit = call_length_scripted(tensor)
print("Scripted: ", length_from_jit)  # Prints 2

call_new_method_scripted = torch.jit.script(call_new_method)  # Throws RuntimeError
</code></pre>
<p><em>Issue 1 root cause</em><br>
The root cause seems to be related to the <a href="https://github.com/pytorch/pytorch/blob/296b45f9d3c3395ebe4aa3bfc6f7b739328cc1f9/torch/jit/annotations.py#L303C10-L303C10" rel="noopener nofollow ugc"><code>is_tensor</code> method</a> in the <code>torch.jit.annotations</code> source code which assigns the standard <code>TensorType</code> to subclasses of <code>torch.Tensor</code> ignoring any modified/custom methods or attributes.</p>
<p><strong>Illustration of issue 2</strong><br>
The script below implements a Tensor subclass with a <code>__torch_function__</code> that overrides only <code>torch.nn.functional.linear</code> by adding a large bias to the output. It then feeds such a tensor through a model consisting of a linear project and a sigmoid activation in eager and scripted mode.</p>
<p>The output shows how scripted mode calls <code>__torch_function__</code> with the outer <code>forward</code> method of the scripted Model and therefore executes the regular <code>torch.nn.functional.linear</code> instead of our override. That is, <code>__torch_function__</code> is not recorded as part of the scripting and custom behaviour gets ignored.</p>
<pre><code class="lang-auto">=== Executing eager model ===
Calling __torch_function__ for &lt;built-in function linear&gt; with name 'linear'
Calling function override for &lt;built-in function linear&gt;
Calling linear override
Calling __torch_function__ for &lt;method 'sum' of 'torch._C._TensorBase' objects&gt; with name 'sum'
Calling __torch_function__ for &lt;built-in method sigmoid of type object at 0x7f192bcfd540&gt; with name 'sigmoid'
Eager:  TensorSubClass([1., 1.], grad_fn=&lt;AliasBackward0&gt;)

=== Executing scripted model ===
Calling __torch_function__ for &lt;torch.ScriptMethod object at 0x7f1828d36de0&gt; with name 'forward'
Scripted:  TensorSubClass([0.5440, 0.4919], grad_fn=&lt;AliasBackward0&gt;)```
</code></pre>
<p><em>Issue 2 script:</em></p>
<pre><code class="lang-python">from typing import Any, Union

import torch


def linear_override(input, weight, bias=None):
    print("Calling linear override")
    x = torch.nn.functional.linear(input, weight, bias)
    return x + 10.0


OVERRIDES = {
    torch.nn.functional.linear: linear_override,
}
    

class TensorSubClass(torch.Tensor):
    def __init__(self, tensor: torch.Tensor):
        super().__init__()
        self.tensor = tensor

    @classmethod
    def to_tensor(cls, input: Any) -&gt; Union[torch.Tensor, Any]:
        return input.tensor if isinstance(input, TensorSubClass) else input

    @classmethod
    def __torch_function__(cls, func, types, args=(), kwargs=None):
        if func is cls.__repr__:
            return super().__torch_function__(func, types, args, kwargs)

        print(f"Calling __torch_function__ for {func} with name '{func.__name__}'")
        if kwargs is None:
            kwargs = {}

        if func not in OVERRIDES:
            return super().__torch_function__(func, types, args, kwargs)

        print(f"Calling function override for {func}")

        # To avoid infinite recursion we first convert all tensors to torch.Tensor. 
        args = torch.utils._pytree.tree_map(cls.to_tensor, args)
        kwargs = torch.utils._pytree.tree_map(cls.to_tensor, kwargs)

        # Call override, convert tensors to TensorSubClass and reconstruct pytrees.
        out = OVERRIDES[func](*args, **kwargs)
        return torch.utils._pytree.tree_map(cls, out)


class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(4, 2)
        self.sigmoid = torch.nn.Sigmoid()

    def forward(self, x: TensorSubClass):
        return self.sigmoid(self.linear(x).sum(-1))


torch.manual_seed(0)

model = Model()
scripted_model = torch.jit.script(model)

x = TensorSubClass(torch.rand((2, 4)))

print("\n=== Executing eager model ===")
out_from_eager = model(x)
# Prints `TensorSubClass([1., 1.], grad_fn=&lt;AliasBackward0&gt;)`
print("Eager: ", out_from_eager)

print("\n=== Executing scripted model ===")
out_from_scripted = scripted_model(x)
# Prints `TensorSubClass([0.5440, 0.4919], grad_fn=&lt;AliasBackward0&gt;)`
print("Scripted: ", out_from_scripted, "\n")
</code></pre>
<p><strong>Potential solution for issue 2</strong><br>
A potential solution to issue 2 seems to be extending PyTorch via <a href="https://dev-discuss.pytorch.org/t/what-and-why-is-torch-dispatch/557" rel="noopener nofollow ugc"><code>__torch_dispatch__</code></a> instead of <code>__torch_function__</code>. It intercepts <code>aten:---</code> operators instead of torch function calls but doesn’t seem quite as well-documented. Importantly, most of these operators are still dispatched even after scripting (those that are changed to other <code>aten:---</code> operators could then also be overridden – we know the full set of <code>aten:---</code> operators a priori). However, this is quite low-level so some extensions could be hard or impossible to make, or just better suited for <code>__torch_function__</code>. Also, it doesn’t solve the issue 1, and it’s unclear to me whether this would work in practice, and whether it would allow ONNX export.</p>