RuntimeError: Unknown type name 'nn.Module' - How to jit a function whose arguments contain an nn.Module?

AlphaBetaGamma96: <p>Hi All,</p>
<p>I was just wondering how <code>torch.jit.script</code> can be used on functions that take <code>nn.Module</code> as an argument as well as <code>Tensor</code>?</p>
<p>For example, I have an example script below which takes the laplacian of a given function, and the <code>laplacian_jit</code> function takes in 2 arguments; the function, <code>net</code>, and the input <code>x</code> (of which we are taking the laplacian). However, when running this it fails with the following error,</p>
<pre><code class="lang-auto">RuntimeError: 
Unknown type name 'nn.Module':
  File "test_jit_func_with_module.py", line 31
@torch.jit.script
def laplacian_jit(net: nn.Module, xs: Tensor):
                       ~~~~~~~~~ &lt;--- HERE
  xis = [xi.requires_grad_() for xi in xs.flatten(start_dim=1).t()]
  xs_flat = torch.stack(xis, dim=1)

</code></pre>
<p>It seems that JIT doesn’t support passing <code>nn.Module</code> as an argument type? Is there a way to define the type such that I can pass an <code>nn.Module</code> type object into the jitted-function?</p>
<p>The example script is below,</p>
<p>Any help will be greatly appreciated!</p>
<p>Thank you!</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

from typing import List, Optional
from torch import Tensor

class Net(nn.Module):
  
  def __init__(self):
    super(Net, self).__init__()
    
  def forward(self, x):
    return x.pow(2).sum(dim=-1)
    
net = Net()

@torch.jit.script
def sumit(inp: List[Optional[torch.Tensor]]):
  elt = inp[0]
  if elt is None:
      raise RuntimeError("blah")
  base = elt
  for i in range(1, len(inp)):
    next_elt = inp[i]
    if next_elt is None:
        raise RuntimeError("blah")
    base = base + next_elt
  return base

@torch.jit.script
def laplacian_jit(net: nn.Module, xs: Tensor):
  xis = [xi.requires_grad_() for xi in xs.flatten(start_dim=1).t()]
  xs_flat = torch.stack(xis, dim=1)
  ys = net(xs_flat.view_as(xs))

  ones = torch.ones_like(ys)
  grad_outputs = torch.jit.annotate(List[Optional[Tensor]], [])
  grad_outputs.append(ones)
  result = torch.autograd.grad([ys], [xs_flat], grad_outputs, retain_graph=True, create_graph=True)
  dy_dxs = result[0]
  if dy_dxs is None:
      raise RuntimeError("blah")

  generator_as_list = [dy_dxs[..., i] for i in range(len(xis))]
  lap_ys_components = [torch.autograd.grad([dy_dxi], [xi], grad_outputs, retain_graph=True, create_graph=False)[0] \
                          for xi, dy_dxi in zip(xis,generator_as_list)]

  lap_ys = sumit(lap_ys_components)

  return lap_ys

x = torch.randn(4096,2)

laplacian_jit(x)
</code></pre>

William: <p>Did you find any answer ? I have a similar problem with a model which can take an other model in the <em>init</em> (try to export Segment Anything Model)</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/william">@William</a>,</p>
<p>I never found a way to jit an <code>nn.Module</code>, but you could perhaps try using <code>functools.partial</code> (docs: <a href="https://docs.python.org/3/library/functools.html#functools.partial" rel="noopener nofollow ugc">here</a>) to create a partial object that treats the <code>nn.Module</code> a static object that shouldn’t be jitted? Perhaps that might be an idea to try?</p>

William: <p>Hi <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a>,<br>
Thank you very much for you answer and your ideas ! You are right, I will see what I can do with it !</p>
<p>Best regards,<br>
William</p>