[jit] Save source maps for imported code error reporting #21183
Description
We should copy JavaScript and optionally save source maps along with the serialized model so that errors on imported code don't show their highlight on the Python-printed code
class M(torch.jit.ScriptModule):
def __init__(self):
super(M, self).__init__()
self.conv = nn.Conv2d(5, 5, 2)
@torch.jit.script_method
def forward(self, x):
return self.conv(x)
M().save("m.pt")
loaded = torch.jit.load("m.pt")
loaded(torch.ones(20, 20))prints
RuntimeError:
Expected 4-dimensional input for 4-dimensional weight 5 5, but got 2-dimensional input of size [20, 20] instead:
operation failed in interpreter:
ret10 = x
ret5 = ret10
ret0 = ret5
ret = ret0
_122 = _0.weight
_123 = _0.bias
_124 = [0, 0]
_1 = torch.conv2d(ret, _122, _123, [1, 1], _124, [1, 1], 1)
else:
_1 = torch.conv2d(x, _0.weight, _0.bias, [1, 1], [0, 0], [1, 1], 1)
~~~~~~~~~~~~ <--- HERE
return _1
but should print
Expected 4-dimensional input for 4-dimensional weight 5 5, but got 2-dimensional input of size [20, 20] instead:
operation failed in interpreter:
@weak_script_method
def forward(self, input):
if self.padding_mode == 'circular':
expanded_padding = ((self.padding[1] + 1) // 2, self.padding[1] // 2,
(self.padding[0] + 1) // 2, self.padding[0] // 2)
return F.conv2d(F.pad(input, expanded_padding, mode='circular'),
self.weight, self.bias, self.stride,
_pair(0), self.dilation, self.groups)
return F.conv2d(input, self.weight, self.bias, self.stride,
~~~~~~~~ <--- HERE
self.padding, self.dilation, self.groups)