torch.nn.ZeroPad2d is not scriptable

[BowenBao]

🐛 Bug

Runtime error thrown when trying to script model that calls nn.ZeroPad2d.

To Reproduce

Steps to reproduce the behavior:

1. Run the below script in python.

import torch
from torch import nn

class Pad(torch.nn.Module):
    def forward(self, x):
        pad_op =  nn.ZeroPad2d(padding=(10, 20, 0, 0))
        return pad_op(x)

m = torch.jit.script(Pad())

2. Observe error

RuntimeError:
Unknown type name '_size_4_t':
  File "/home/bowbao/repos/pytorch/torch/nn/modules/padding.py", line 453
    def __init__(self, padding: _size_4_t) -> None:
                                ~~~~~~~~~ <--- HERE
        super(ZeroPad2d, self).__init__(padding, 0.)
'ZeroPad2d.__init__' is being compiled since it was called from '__torch__.torch.nn.modules.padding.ZeroPad2d'
  File "repro_1p_rnn_fairseq.py", line 6
    def forward(self, x):
        pad_op =  nn.ZeroPad2d(padding=(10, 20, 0, 0))
                  ~~~~~~~~~~~~ <--- HERE
        return pad_op(x)
'__torch__.torch.nn.modules.padding.ZeroPad2d' is being compiled since it was called from 'Pad.forward'
  File "repro_1p_rnn_fairseq.py", line 6
    def forward(self, x):
        pad_op =  nn.ZeroPad2d(padding=(10, 20, 0, 0))
        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
        return pad_op(x)

Expected behavior

Create scripted module correctly.

cc @gmagogsfm

[ssnl]

Just curious- do we support such things given that the correct thing is to either create the module (and assign as attr) in __init__ or use the functional form?

[wanchaol]

@malfet since you are adding inline type annotations for nn.Module inference, this might be a case of regression, can you take a look?

[malfet]

The problem originates from the fact that torch.scipt does not support Unions, example below

size_2_t = Union[float, Tuple[float, float]]
def foo(x: size_2_t) -> float: return x[0] + 1.0
f=torch.jit.script(foo)

results in the same non-descriptive Unknown type name 'size_2_t':, although if one to remove the Union it will work.
Although adding a generic Union support sounds like a complicated task, I wonder if adding internal list_or_scalar_type would be an acceptable workaround, because it's a very common pattern in PyTorch

[eellison]

I think this is what BroadcastingList is equivalent to

[malfet]

@eellison, I see, so this should work:

from torch._jit_internal import BroadcastingList2
def foo(x: BroadcastingList2[float]) -> float: return x[0] + x[1]

print(torch.jit.script(foo)(3))