Torch.jit.script throws error related to pytorch code

username: rwesterman
cooked: <p>I’m seeing the following error when trying to compile a model to torchscript using <code>torch.jit.script</code>:</p>
<pre><code class="lang-auto">Traceback (most recent call last):                                                                                                     
  File "espnet_to_jit.py", line 31, in &lt;module&gt;
    script_model = torch.jit.script(asr_model)                                                                                         
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/jit/_script.py", line 943, in script
    obj, torch.jit._recursive.infer_methods_to_compile                                                                                                   
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/jit/_recursive.py", line 391, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)                                                               
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/jit/_recursive.py", line 448, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)                                                    
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/jit/_script.py", line 391, in _construct
    init_fn(script_module)                                                                                                                               
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/jit/_recursive.py", line 428, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)                                                      
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/jit/_recursive.py", line 448, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)                                                    
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/jit/_script.py", line 391, in _construct
    init_fn(script_module)                                                                                                                               
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/jit/_recursive.py", line 428, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)                                                      
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/jit/_recursive.py", line 448, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)                                                    
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/jit/_script.py", line 391, in _construct
    init_fn(script_module)                                                                                                                               
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/jit/_recursive.py", line 428, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)                                                                                       
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/jit/_recursive.py", line 448, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)                                                                   
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/jit/_script.py", line 391, in _construct
    init_fn(script_module)                                                                                                                             
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/jit/_recursive.py", line 428, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)                                                              
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/jit/_recursive.py", line 452, in create_script_module_impl
    create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)                                               
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/jit/_recursive.py", line 335, in create_methods_and_properties_from_stubs
    concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)                
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/jit/_script.py", line 1106, in _recursive_compile_class
    _compile_and_register_class(obj, rcb, _qual_name)                                                                                         
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/jit/_script.py", line 65, in _compile_and_register_class
    ast = get_jit_class_def(obj, obj.__name__)                                              
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/jit/frontend.py", line 176, in get_jit_class_def
    is_classmethod=is_classmethod(method[1])) for method in methods]          
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/jit/frontend.py", line 176, in &lt;listcomp&gt;
    is_classmethod=is_classmethod(method[1])) for method in methods]                                       
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/jit/frontend.py", line 248, in get_jit_def
    type_line = torch.jit.annotations.get_type_line(source)                                                    
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/jit/annotations.py", line 205, in get_type_line
    "\n(See PEP 484 https://www.python.org/dev/peps/pep-0484/#suggested-syntax-for-python-2-7-and-straddling-code)")  # noqa
RuntimeError: Return type line '# type: (...) -&gt; ...' not found on multiline type annotation
for type lines:                                                               
            normalized_shape = (normalized_shape,)  # type: ignore[assignment]   
        self.normalized_shape = tuple(normalized_shape)  # type: ignore[arg-type]
(See PEP 484 https://www.python.org/dev/peps/pep-0484/#suggested-syntax-for-python-2-7-and-straddling-code)
</code></pre>
<p>I added to the code in <code>annotations.py</code> to print out the source of the error as well, and it shows:</p>
<pre><code class="lang-auto">Source is:
    def __init__(self, normalized_shape: _shape_t, eps: float = 1e-5, elementwise_affine: bool = True) -&gt; None:
        super(LayerNorm, self).__init__()
        if isinstance(normalized_shape, numbers.Integral):
            # mypy error: incompatible types in assignment
            normalized_shape = (normalized_shape,)  # type: ignore[assignment]
        self.normalized_shape = tuple(normalized_shape)  # type: ignore[arg-type]
        self.eps = eps
        self.elementwise_affine = elementwise_affine
        if self.elementwise_affine:
            self.weight = Parameter(torch.Tensor(*self.normalized_shape))
            self.bias = Parameter(torch.Tensor(*self.normalized_shape))
        else:
            self.register_parameter('weight', None)
            self.register_parameter('bias', None)
        self.reset_parameters()
</code></pre>
<p>I tracked down the source of this code to <code>torch/nn/modules/normalization.py</code>.</p><aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/normalization.py#L148" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/normalization.py#L148" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/master/torch/nn/modules/normalization.py#L148</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="138" style="counter-reset: li-counter 137 ;">
<li>    &gt;&gt;&gt; # Normalize over last dimension of size 10</li>
<li>    &gt;&gt;&gt; m = nn.LayerNorm(10)</li>
<li>    &gt;&gt;&gt; # Activating the module</li>
<li>    &gt;&gt;&gt; output = m(input)</li>
<li>"""</li>
<li>__constants__ = ['normalized_shape', 'eps', 'elementwise_affine']</li>
<li>normalized_shape: Tuple[int, ...]</li>
<li>eps: float</li>
<li>elementwise_affine: bool</li>
<li>
</li>
<li class="selected">def __init__(self, normalized_shape: _shape_t, eps: float = 1e-5, elementwise_affine: bool = True,</li>
<li>             device=None, dtype=None) -&gt; None:</li>
<li>    factory_kwargs = {'device': device, 'dtype': dtype}</li>
<li>    super(LayerNorm, self).__init__()</li>
<li>    if isinstance(normalized_shape, numbers.Integral):</li>
<li>        # mypy error: incompatible types in assignment</li>
<li>        normalized_shape = (normalized_shape,)  # type: ignore[assignment]</li>
<li>    self.normalized_shape = tuple(normalized_shape)  # type: ignore[arg-type]</li>
<li>    self.eps = eps</li>
<li>    self.elementwise_affine = elementwise_affine</li>
<li>    if self.elementwise_affine:</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>Why would <code>torch.jit.script</code> be tracing into the pytorch code, and how can I fix this error?</p>

username: rwesterman
cooked: <p>I was able to move past this issue by removing the <code># type: ignore[]</code> comments, but am immediately presented with another error in compiling pytorch code:</p>
<pre><code class="lang-auto">torch.jit.frontend.NotSupportedError: keyword-arg expansion is not supported:
  File "/shared/workspaces/rwesterman/miniconda3/envs/espnet/lib/python3.7/site-packages/torch/nn/modules/normalization.py", line 175
    def extra_repr(self) -&gt; str:
        return '{normalized_shape}, eps={eps}, ' \
            'elementwise_affine={elementwise_affine}'.format(**self.__dict__)
                                                               ~~~~~~~~~~~~~ &lt;--- HERE
'__torch__.torch.nn.modules.normalization.LayerNorm' is being compiled since it was called from 'LayerNorm.forward'
  File "/shared/workspaces/rwesterman/rev-espnet/espnet/nets/pytorch_backend/transformer/layer_norm.py", line 37
        """
        if self.dim == -1:
            return super(torch.nn.LayerNorm, self).forward(x)
                         ~~~~~~~~~~~~~~~~~~ &lt;--- HERE
        return super(torch.nn.LayerNorm, self).forward(x.transpose(1, -1)).transpose(1, -1)
</code></pre>