torchScript failed with @ __matmul__ operator

jiahuil: <p>hi all, I am trying to define a new @ operator in a class, then use it in torchScript model, but it failed.  However when I test on <strong>add</strong> operater, it works. Here is the demo below:</p>
<p>Source code:</p>
<pre><code class="lang-python">import torch
import traceback

class NewTensor:
    def __init__(self, value: torch.Tensor):
        self.value = value  

    def __matmul__(self, other: "NewTensor") -&gt; torch.Tensor:
        return -torch.matmul(self.value, other.value)

    def __add__(self, other: "NewTensor") -&gt; torch.Tensor:
        return -torch.add(self.value, other.value)


class mul_model:
    def __init__(self) -&gt; None:
        pass

    def forward(self, x: NewTensor, y: NewTensor):
        return x @ y


class add_model:
    def __init__(self) -&gt; None:
        pass

    def forward(self, x: NewTensor, y: NewTensor):
        return x + y


model_add = add_model()
try:
    a_s = torch.jit.script(model_add)
    print("add model success")
except Exception as e:
    print("add model fail")
    traceback.print_exc()

model_mul = mul_model()
try:
    a_s = torch.jit.script(model_mul)
    print("mul model success")
except:
    print("mul model fail")
    traceback.print_exc()
</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">add model success
mul model fail
Traceback (most recent call last):
  File "test.py", line 42, in &lt;module&gt;
    a_s = torch.jit.script(model_mul)
  File "/opt/conda/lib/python3.8/site-packages/torch/jit/_script.py", line 1351, in script
    return torch.jit._recursive.create_script_class(obj)
  File "/opt/conda/lib/python3.8/site-packages/torch/jit/_recursive.py", line 428, in create_script_class
    _compile_and_register_class(type(obj), rcb, qualified_class_name)
  File "/opt/conda/lib/python3.8/site-packages/torch/jit/_recursive.py", line 46, in _compile_and_register_class
    script_class = torch._C._jit_script_class_compile(qualified_name, ast, defaults, rcb)
RuntimeError:
Arguments for call are not valid.
The following variants are available:

  aten::matmul(Tensor self, Tensor other) -&gt; Tensor:
  Expected a value of type 'Tensor' for argument 'self' but instead found type '__torch__.NewTensor (of Python compilation unit at: 0x55951834bdb0)'.

  aten::matmul.out(Tensor self, Tensor other, *, Tensor(a!) out) -&gt; Tensor(a!):
  Expected a value of type 'Tensor' for argument 'self' but instead found type '__torch__.NewTensor (of Python compilation unit at: 0x55951834bdb0)'.

The original call is:
  File "test.py", line 21
    def forward(self, x: NewTensor, y: NewTensor):
        return x @ y
               ~~~~~ &lt;--- HERE

</code></pre>
<p>Anyone come across similar issue?</p>
<p>Thanks</p>

ptrblck: <p>I’m unsure if TorchScript is supposed to support custom classes, so could you check if creating a custom Tensor class via <code>__torch_function__</code> would work instead as described <a href="https://pytorch.org/docs/stable/notes/extending.html#extending-torch-with-a-tensor-like-type">here</a>?</p>

jiahuil: <p>thanks, I will try it. <img src="https://discuss.pytorch.org/images/emoji/apple/grinning.png?v=12" title=":grinning:" class="emoji" alt=":grinning:" loading="lazy" width="20" height="20"></p>

jiahuil: <p>I try some thing like below, but still failed.</p>
<pre><code class="lang-auto">class mul_model(torch.nn.Module):
    def __init__(self) -&gt; None:
        super().__init__()

    def forward(self, x: NewTensor, y: NewTensor):
        return x @ y


class add_model(torch.nn.Module):
    def __init__(self) -&gt; None:
        super().__init__()

    def forward(self, x: NewTensor, y: NewTensor):
        return x + y
</code></pre>
<pre><code class="lang-auto">add model success
Traceback (most recent call last):
  File "test.py", line 97, in &lt;module&gt;
    a_s = torch.jit.script(model_mul)
  File "/opt/conda/lib/python3.8/site-packages/torch/jit/_script.py", line 1286, in script
    return torch.jit._recursive.create_script_module(
  File "/opt/conda/lib/python3.8/site-packages/torch/jit/_recursive.py", line 460, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/opt/conda/lib/python3.8/site-packages/torch/jit/_recursive.py", line 526, in create_script_module_impl
    create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
  File "/opt/conda/lib/python3.8/site-packages/torch/jit/_recursive.py", line 377, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
RuntimeError:
Arguments for call are not valid.
The following variants are available:

  aten::matmul(Tensor self, Tensor other) -&gt; Tensor:
  Expected a value of type 'Tensor' for argument 'self' but instead found type '__torch__.NewTensor (of Python compilation unit at: 0x55af73107d90)'.

  aten::matmul.out(Tensor self, Tensor other, *, Tensor(a!) out) -&gt; Tensor(a!):
  Expected a value of type 'Tensor' for argument 'self' but instead found type '__torch__.NewTensor (of Python compilation unit at: 0x55af73107d90)'.

The original call is:
  File "test.py", line 21
    def forward(self, x: NewTensor, y: NewTensor):
        return x @ y
               ~~~~~ &lt;--- HERE
</code></pre>
<p>It is weird that torchScipt fails on matmul but works on add operator.</p>

ptrblck: <p>Did you try to add this operation via a <code>__torch_function__</code>?</p>

jiahuil: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a> ,</p>
<p>I tried <code>__torch_function__</code> but it cannot be scripted by TorchScript because TorchScript analyzes the code statically.</p>
<p>Basically, I want to define a custom class <code>NewTensor</code> that can be scripted by TorchScript following the practice of <a href="https://pytorch.org/docs/stable/jit_language_reference.html#id2" rel="noopener nofollow ugc">TorchScript custom classes</a>. The NewTensor has both <code>__add__</code> and <code>__matmul__</code> functions to support + and @ operators. However, we found that  TorchScript works well when scripting <code>+</code> but fails on <code>@</code>. I am wondering whether this is a bug of TorchScript.</p>
<pre><code class="lang-auto">@torch.jit.script
class NewTensor:
    def __init__(self, value: torch.Tensor):
        self.value = value  

    def __matmul__(self, other: "NewTensor") -&gt; torch.Tensor:
        return -torch.matmul(self.value, other.value)

    def __add__(self, other: "NewTensor") -&gt; torch.Tensor:
        return -torch.add(self.value, other.value)
</code></pre>
<p>In the code above, if I use <code>x + y</code> where <code>x</code> and <code>y</code> are <code>NewTensor</code>, TorchScript can recognize the operand as <code>NewTensor.__add__(x, y)</code>. But if I use <code>x @ y</code>, TorchScript will raise an error as above.</p>