Issue 1:
Trying to use pack_padded_sequence through a torch.jit.ScriptModule reports the following prototype which seems false ? pack_padded_sequence's prototype is (Tensor, Tensor, bool, bool).

RuntimeError:
for operator (Tensor 0, Tensor 1, Tensor 2, Tensor 3) -> Tensor:
argument 0 not provided.

When trying to use the following module:

from typing import List

import torch
from torch.jit import ScriptModule, script_method, annotate
from torch.nn import Dropout, ModuleList, ParameterList, Parameter, GRU
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

__author__ = 'Morgan Funtowicz'

class StackedBiRNN(ScriptModule):
"""
Stacked Bi-directional RNNs.
Differs from standard PyTorch library in that it has the option to save
and concat the hidden states between layers. (i.e. the output hidden size
for each sequence input is num_layers * hidden_size).
"""

__constants__ = ['padding', 'dropout_rate', 'num_layers', 'concat_layers', '_rnns']

def __init__(self, input_size, hidden_size, num_layers,
dropout_p=0., rnn_type=GRU,
concat=False, padding=False):
super(StackedBiRNN, self).__init__()

self.padding = padding
self.dropout_rate = dropout_p
self.num_layers = num_layers
self.concat_layers = concat

rnns = []
for i in range(num_layers):
input_size = input_size if i == 0 else 2 * hidden_size

rnns.append(rnn_type(input_size, hidden_size, batch_first=True, bidirectional=True))

self._rnns = ModuleList(rnns)
self._dropout = Dropout(dropout_p)

@property
def hidden_size(self):
return self._rnns[0].hidden_size

@script_method
def forward(self, x: torch.Tensor, lengths: torch.Tensor):

# Buffer storing RNN layers output
outputs = annotate(List[torch.Tensor], [x])

# Iterate through each RNN layer
for rnn in self._rnns:

# Retrieve the latest output (output from previous layer)
rnn_input = outputs[-1]

# Apply dropout to input
if self.dropout_rate > 0:
rnn_input = self._dropout(rnn_input)

rnn_input = pack_padded_sequence(input=rnn_input, lengths=lengths, batch_first=True, enforce_sorted=False)

# Go through the RNN and get output
rnn.flatten_parameters()
out = rnn(rnn_input)[0]
outputs += [pad_packed_sequence(sequence=out, batch_first=True)[0]]

# Concat hidden layers or take final
if self.concat_layers:
return torch.cat(outputs[1:], dim=-1)
else:
return outputs[-1]

Issue 2:
🚀 Feature

Add TorchScript support for torch.nn.utils.rnn.pack_padded_sequence() and torch.nn.utils.rnn.pad_packed_sequence().

Motivation

Motivated by https://discuss.pytorch.org/t/torch-jit-script-for-torch-nn-utils-rnn-pack-padded-sequence/30668, this appears to be a desired feature given that tracing is usually not an option when these functions are used (for variable length sequences)

Pitch

Right now, the following code:

from torch import nn
import torch.jit

class MyLSTM(torch.jit.ScriptModule):
def __init__(self):

super(MyLSTM, self).__init__()

self.rnn = nn.LSTM(input_size=300,
hidden_size=32,
batch_first=True,
num_layers=1,
bidirectional=False)

@torch.jit.script_method
def forward(self, x_in, x_lengths):
max_length = x_lengths[0]
x_in_shortened = x_in[:,:max_length]

x_packed = nn.utils.rnn.pack_padded_sequence(x_in_shortened,
x_lengths,
True, # batch_first
True # enforce_sorted
)

lstm_outputs = self.rnn(x_packed)
sequence_outputs = nn.utils.rnn.pad_packed_sequence(lstm_outputs[0],
batch_first=True,
padding_value=0.0,
total_length=max_length)

lstm_out = sequence_outputs[0]
out_seq_length = sequence_outputs[1]

return lstm_out

# Instantiation:
my_lstm = MyLSTM()

Fails with:

for operator (Tensor 0, Tensor 1, Tensor 2, Tensor 3) -> Tensor:
expected a value of type Tensor for argument '2' but found bool
@torch.jit.script_method
def forward(self, x_in, x_lengths):
max_length = x_lengths[0]
x_in_shortened = x_in[:,:max_length]

x_packed = nn.utils.rnn.pack_padded_sequence(x_in_shortened,
x_lengths,
True, # batch_first
~~~~ <--- HERE
True # enforce_sorted
)

So it seems that it expects Tensors for the arguments batch_first and enforce_sorted instead of Python bools. Okay, if we change them to:

x_packed = nn.utils.rnn.pack_padded_sequence(x_in_shortened,
x_lengths,
torch.ones(1).byte().squeeze(), # batch_first
torch.ones(1).byte().squeeze() # enforce_sorted
)

Insantiation works, but it fails when the forward method is applied:

my_lstm(torch.Tensor([[3,1,2,3,4]]), torch.Tensor([5]).long())

# Fails with:
TypeError: _pack_padded_sequence(): argument 'batch_first' (position 3) must be bool, not Tensor

So currently, either instantiation or forward pass is always going to fail. I believe that this is just argument parsing, and that the "hard" part of this issue is to actually make the functions compatible with TorchScript...

BTW Thank you for PyTorch overall. ¡It's awesome!