Callable support in Torchscript IR

sambhav_jain: <p>Hi PyTorch. How can I get <code>jit.script</code> to represent a Callable in the IR?</p>
<p>Specifically, in the example shown below, we use a <code>Tensor.map_</code> to map the provided functor (in this case to compute euclidean distance between two elements) pointwise on two tensors.<br>
Basically we’re looking for a way to represent <code>aten::map_</code> in the JIT IR, and have the interpreter  receive a function ptr to apply for each element in the operand tensors.</p>
<pre><code class="lang-auto"># test1.py

import torch
from typing import Callable

class Map(torch.nn.Module):
    def forward(self, x: torch.Tensor, y: torch.Tensor, functor: Callable):
        x.map_(y, functor)
        return x

def euclidean_dist(x: float, y: float) -&gt; float:
    return (x - y) * (x + y)

x = torch.tensor([1.0, 2.0, 3.0, 6.0])
y = torch.tensor([5.0, -1.0, 4.0, 3.5])

model = Map()
print(model(x, y, euclidean_dist))

script = torch.jit.script(model)
</code></pre>
<p>This expectedly fails (since <a href="https://pytorch.org/docs/stable/jit_language_reference.html#unsupported-typing-constructs" rel="noopener nofollow ugc">Callable is not supported</a>):</p>
<pre><code class="lang-auto">RuntimeError: 
Unknown type name 'Callable':
  File "test1.py", line 5
    def forward(self, x: torch.Tensor, y: torch.Tensor, functor: Callable):
                                                                 ~~~~~~~~ &lt;--- HERE
        x.map_(y, functor)
        return x
</code></pre>
<p>Tried to use <code>torch.jit.interface</code> to specify the interface contract for the Callable, but I get a different error now:</p>
<pre><code class="lang-auto"># test2.py

import torch

@torch.jit.interface
class FunctorInterface(object):
    def forward(self, x: float, y: float) -&gt; float:
        pass

class Map(torch.nn.Module):
    def forward(self, x: torch.Tensor, y: torch.Tensor, functor: FunctorInterface):
        x.map_(y, functor.forward)
        return x

class EuclideanDist(torch.nn.Module):
    def forward(self, x: float, y: float) -&gt; float:
        return (x - y) * (x + y)

x = torch.tensor([1.0, 2.0, 3.0, 6.0])
y = torch.tensor([5.0, -1.0, 4.0, 3.5])

model = Map()
print(model(x, y, EuclideanDist()))

script = torch.jit.script(model)
</code></pre>
<p>Error:</p>
<pre><code class="lang-auto">RuntimeError: 
'Tensor' object has no attribute or method 'map_'.:
  File "test2.py", line 10
    def forward(self, x: torch.Tensor, y: torch.Tensor, functor: FunctorInterface):
        x.map_(y, functor.forward)
        ~~~~~~ &lt;--- HERE
        return x
</code></pre>
<p>Modifying test1 slightly, say we could specify the concrete lambda at model instantiation. For example:</p>
<pre><code class="lang-auto"># test3.py

import torch
from typing import Callable

class Map(torch.nn.Module):
    def __init__(self, functor: Callable):
        super().__init__()
        self.functor = functor

    def forward(self, x: torch.Tensor, y: torch.Tensor):
        x.map_(y, self.functor)
        return x

def euclidean_dist(x: float, y: float) -&gt; float:
    return (x - y) * (x + y)

x = torch.tensor([1.0, 2.0, 3.0, 6.0])
y = torch.tensor([5.0, -1.0, 4.0, 3.5])

model = Map(euclidean_dist)
print(model(x, y))

script = torch.jit.script(model)
</code></pre>
<p>This still fails:</p>
<pre><code class="lang-auto">RuntimeError:
'Tensor' object has no attribute or method 'map_'.:
  File "test.py", line 12
    def forward(self, x: torch.Tensor, y: torch.Tensor):
        x.map_(y, self.functor)
        ~~~~~~ &lt;--- HERE
        return x
</code></pre>