<p dir="auto">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).</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="RuntimeError: 
for operator (Tensor 0, Tensor 1, Tensor 2, Tensor 3) -&gt; Tensor:
argument 0 not provided."><pre class="notranslate"><span class="pl-v">RuntimeError</span>: 
<span class="pl-s1">for</span> <span class="pl-en">operator</span> (<span class="pl-v">Tensor</span> <span class="pl-c1">0</span>, <span class="pl-v">Tensor</span> <span class="pl-c1">1</span>, <span class="pl-v">Tensor</span> <span class="pl-c1">2</span>, <span class="pl-v">Tensor</span> <span class="pl-c1">3</span>) <span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-v">Tensor</span>:
<span class="pl-s1">argument</span> <span class="pl-c1">0</span> <span class="pl-c1">not</span> <span class="pl-s1">provided</span>.</pre></div>
<p dir="auto">When trying to use the following module:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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):
    &quot;&quot;&quot;
    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).
    &quot;&quot;&quot;

    __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 &gt; 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]"><pre class="notranslate"><span class="pl-k">from</span> <span class="pl-s1">typing</span> <span class="pl-k">import</span> <span class="pl-v">List</span>

<span class="pl-k">import</span> <span class="pl-s1">torch</span>
<span class="pl-k">from</span> <span class="pl-s1">torch</span>.<span class="pl-s1">jit</span> <span class="pl-k">import</span> <span class="pl-v">ScriptModule</span>, <span class="pl-s1">script_method</span>, <span class="pl-s1">annotate</span>
<span class="pl-k">from</span> <span class="pl-s1">torch</span>.<span class="pl-s1">nn</span> <span class="pl-k">import</span> <span class="pl-v">Dropout</span>, <span class="pl-v">ModuleList</span>, <span class="pl-v">ParameterList</span>, <span class="pl-v">Parameter</span>, <span class="pl-v">GRU</span>
<span class="pl-k">from</span> <span class="pl-s1">torch</span>.<span class="pl-s1">nn</span>.<span class="pl-s1">utils</span>.<span class="pl-s1">rnn</span> <span class="pl-k">import</span> <span class="pl-s1">pack_padded_sequence</span>, <span class="pl-s1">pad_packed_sequence</span>

<span class="pl-s1">__author__</span> <span class="pl-c1">=</span> <span class="pl-s">'Morgan Funtowicz'</span>


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

    <span class="pl-s1">__constants__</span> <span class="pl-c1">=</span> [<span class="pl-s">'padding'</span>, <span class="pl-s">'dropout_rate'</span>, <span class="pl-s">'num_layers'</span>, <span class="pl-s">'concat_layers'</span>, <span class="pl-s">'_rnns'</span>]

    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">input_size</span>, <span class="pl-s1">hidden_size</span>, <span class="pl-s1">num_layers</span>,
                 <span class="pl-s1">dropout_p</span><span class="pl-c1">=</span><span class="pl-c1">0.</span>, <span class="pl-s1">rnn_type</span><span class="pl-c1">=</span><span class="pl-v">GRU</span>,
                 <span class="pl-s1">concat</span><span class="pl-c1">=</span><span class="pl-c1">False</span>, <span class="pl-s1">padding</span><span class="pl-c1">=</span><span class="pl-c1">False</span>):
        <span class="pl-en">super</span>(<span class="pl-v">StackedBiRNN</span>, <span class="pl-s1">self</span>).<span class="pl-en">__init__</span>()

        <span class="pl-s1">self</span>.<span class="pl-s1">padding</span> <span class="pl-c1">=</span> <span class="pl-s1">padding</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">dropout_rate</span> <span class="pl-c1">=</span> <span class="pl-s1">dropout_p</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">num_layers</span> <span class="pl-c1">=</span> <span class="pl-s1">num_layers</span>
        <span class="pl-s1">self</span>.<span class="pl-s1">concat_layers</span> <span class="pl-c1">=</span> <span class="pl-s1">concat</span>

        <span class="pl-s1">rnns</span> <span class="pl-c1">=</span> []
        <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">num_layers</span>):
            <span class="pl-s1">input_size</span> <span class="pl-c1">=</span> <span class="pl-s1">input_size</span> <span class="pl-k">if</span> <span class="pl-s1">i</span> <span class="pl-c1">==</span> <span class="pl-c1">0</span> <span class="pl-k">else</span> <span class="pl-c1">2</span> <span class="pl-c1">*</span> <span class="pl-s1">hidden_size</span>

            <span class="pl-s1">rnns</span>.<span class="pl-en">append</span>(<span class="pl-en">rnn_type</span>(<span class="pl-s1">input_size</span>, <span class="pl-s1">hidden_size</span>, <span class="pl-s1">batch_first</span><span class="pl-c1">=</span><span class="pl-c1">True</span>, <span class="pl-s1">bidirectional</span><span class="pl-c1">=</span><span class="pl-c1">True</span>))

        <span class="pl-s1">self</span>.<span class="pl-s1">_rnns</span> <span class="pl-c1">=</span> <span class="pl-v">ModuleList</span>(<span class="pl-s1">rnns</span>)
        <span class="pl-s1">self</span>.<span class="pl-s1">_dropout</span> <span class="pl-c1">=</span> <span class="pl-v">Dropout</span>(<span class="pl-s1">dropout_p</span>)

    <span class="pl-en">@<span class="pl-s1">property</span></span>
    <span class="pl-k">def</span> <span class="pl-en">hidden_size</span>(<span class="pl-s1">self</span>):
        <span class="pl-k">return</span> <span class="pl-s1">self</span>.<span class="pl-s1">_rnns</span>[<span class="pl-c1">0</span>].<span class="pl-s1">hidden_size</span>

    <span class="pl-en">@<span class="pl-s1">script_method</span></span>
    <span class="pl-k">def</span> <span class="pl-en">forward</span>(<span class="pl-s1">self</span>, <span class="pl-s1">x</span>: <span class="pl-s1">torch</span>.<span class="pl-v">Tensor</span>, <span class="pl-s1">lengths</span>: <span class="pl-s1">torch</span>.<span class="pl-v">Tensor</span>):

        <span class="pl-c"># Buffer storing RNN layers output</span>
        <span class="pl-s1">outputs</span> <span class="pl-c1">=</span> <span class="pl-en">annotate</span>(<span class="pl-v">List</span>[<span class="pl-s1">torch</span>.<span class="pl-v">Tensor</span>], [<span class="pl-s1">x</span>])

        <span class="pl-c"># Iterate through each RNN layer</span>
        <span class="pl-k">for</span> <span class="pl-s1">rnn</span> <span class="pl-c1">in</span> <span class="pl-s1">self</span>.<span class="pl-s1">_rnns</span>:

            <span class="pl-c"># Retrieve the latest output (output from previous layer)</span>
            <span class="pl-s1">rnn_input</span> <span class="pl-c1">=</span> <span class="pl-s1">outputs</span>[<span class="pl-c1">-</span><span class="pl-c1">1</span>]

            <span class="pl-c"># Apply dropout to input</span>
            <span class="pl-k">if</span> <span class="pl-s1">self</span>.<span class="pl-s1">dropout_rate</span> <span class="pl-c1">&gt;</span> <span class="pl-c1">0</span>:
                <span class="pl-s1">rnn_input</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">_dropout</span>(<span class="pl-s1">rnn_input</span>)

            <span class="pl-s1">rnn_input</span> <span class="pl-c1">=</span> <span class="pl-en">pack_padded_sequence</span>(<span class="pl-s1">input</span><span class="pl-c1">=</span><span class="pl-s1">rnn_input</span>, <span class="pl-s1">lengths</span><span class="pl-c1">=</span><span class="pl-s1">lengths</span>, <span class="pl-s1">batch_first</span><span class="pl-c1">=</span><span class="pl-c1">True</span>, <span class="pl-s1">enforce_sorted</span><span class="pl-c1">=</span><span class="pl-c1">False</span>)

            <span class="pl-c"># Go through the RNN and get output</span>
            <span class="pl-s1">rnn</span>.<span class="pl-en">flatten_parameters</span>()
            <span class="pl-s1">out</span> <span class="pl-c1">=</span> <span class="pl-en">rnn</span>(<span class="pl-s1">rnn_input</span>)[<span class="pl-c1">0</span>]
            <span class="pl-s1">outputs</span> <span class="pl-c1">+=</span> [<span class="pl-en">pad_packed_sequence</span>(<span class="pl-s1">sequence</span><span class="pl-c1">=</span><span class="pl-s1">out</span>, <span class="pl-s1">batch_first</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)[<span class="pl-c1">0</span>]]

        <span class="pl-c"># Concat hidden layers or take final</span>
        <span class="pl-k">if</span> <span class="pl-s1">self</span>.<span class="pl-s1">concat_layers</span>:
            <span class="pl-k">return</span> <span class="pl-s1">torch</span>.<span class="pl-en">cat</span>(<span class="pl-s1">outputs</span>[<span class="pl-c1">1</span>:], <span class="pl-s1">dim</span><span class="pl-c1">=</span><span class="pl-c1">-</span><span class="pl-c1">1</span>)
        <span class="pl-k">else</span>:
            <span class="pl-k">return</span> <span class="pl-s1">outputs</span>[<span class="pl-c1">-</span><span class="pl-c1">1</span>]</pre></div>
<h2 dir="auto"><g-emoji class="g-emoji" alias="rocket" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/1f680.png">🚀</g-emoji> Feature</h2>
<p dir="auto">Add TorchScript support for <code class="notranslate">torch.nn.utils.rnn.pack_padded_sequence()</code> and <code class="notranslate">torch.nn.utils.rnn.pad_packed_sequence()</code>.</p>
<h2 dir="auto">Motivation</h2>
<p dir="auto">Motivated by <a href="https://discuss.pytorch.org/t/torch-jit-script-for-torch-nn-utils-rnn-pack-padded-sequence/30668" rel="nofollow">https://discuss.pytorch.org/t/torch-jit-script-for-torch-nn-utils-rnn-pack-padded-sequence/30668</a>, this appears to be a desired feature given that tracing is usually not an option when these functions are used (for variable length sequences)</p>
<h2 dir="auto">Pitch</h2>
<p dir="auto">Right now, the following code:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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()"><pre class="notranslate"><span class="pl-k">from</span> <span class="pl-s1">torch</span> <span class="pl-k">import</span> <span class="pl-s1">nn</span>
<span class="pl-k">import</span> <span class="pl-s1">torch</span>.<span class="pl-s1">jit</span>

<span class="pl-k">class</span> <span class="pl-v">MyLSTM</span>(<span class="pl-s1">torch</span>.<span class="pl-s1">jit</span>.<span class="pl-v">ScriptModule</span>):
    <span class="pl-k">def</span> <span class="pl-en">__init__</span>(<span class="pl-s1">self</span>):

        <span class="pl-en">super</span>(<span class="pl-v">MyLSTM</span>, <span class="pl-s1">self</span>).<span class="pl-en">__init__</span>()
        
        <span class="pl-s1">self</span>.<span class="pl-s1">rnn</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-v">LSTM</span>(<span class="pl-s1">input_size</span><span class="pl-c1">=</span><span class="pl-c1">300</span>,
                           <span class="pl-s1">hidden_size</span><span class="pl-c1">=</span><span class="pl-c1">32</span>,
                           <span class="pl-s1">batch_first</span><span class="pl-c1">=</span><span class="pl-c1">True</span>,
                           <span class="pl-s1">num_layers</span><span class="pl-c1">=</span><span class="pl-c1">1</span>,
                           <span class="pl-s1">bidirectional</span><span class="pl-c1">=</span><span class="pl-c1">False</span>)
    
    <span class="pl-en">@<span class="pl-s1">torch</span>.<span class="pl-s1">jit</span>.<span class="pl-s1">script_method</span></span>
    <span class="pl-k">def</span> <span class="pl-en">forward</span>(<span class="pl-s1">self</span>, <span class="pl-s1">x_in</span>, <span class="pl-s1">x_lengths</span>):
        <span class="pl-s1">max_length</span> <span class="pl-c1">=</span> <span class="pl-s1">x_lengths</span>[<span class="pl-c1">0</span>]
        <span class="pl-s1">x_in_shortened</span> <span class="pl-c1">=</span> <span class="pl-s1">x_in</span>[:,:<span class="pl-s1">max_length</span>]
        
        <span class="pl-s1">x_packed</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-s1">utils</span>.<span class="pl-s1">rnn</span>.<span class="pl-en">pack_padded_sequence</span>(<span class="pl-s1">x_in_shortened</span>, 
                                                     <span class="pl-s1">x_lengths</span>,
                                                     <span class="pl-c1">True</span>, <span class="pl-c"># batch_first</span>
                                                     <span class="pl-c1">True</span> <span class="pl-c"># enforce_sorted</span>
                                                    )
        
        <span class="pl-s1">lstm_outputs</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">rnn</span>(<span class="pl-s1">x_packed</span>)
        <span class="pl-s1">sequence_outputs</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-s1">utils</span>.<span class="pl-s1">rnn</span>.<span class="pl-en">pad_packed_sequence</span>(<span class="pl-s1">lstm_outputs</span>[<span class="pl-c1">0</span>], 
                                                            <span class="pl-s1">batch_first</span><span class="pl-c1">=</span><span class="pl-c1">True</span>, 
                                                            <span class="pl-s1">padding_value</span><span class="pl-c1">=</span><span class="pl-c1">0.0</span>, 
                                                            <span class="pl-s1">total_length</span><span class="pl-c1">=</span><span class="pl-s1">max_length</span>)
        
        <span class="pl-s1">lstm_out</span> <span class="pl-c1">=</span> <span class="pl-s1">sequence_outputs</span>[<span class="pl-c1">0</span>]
        <span class="pl-s1">out_seq_length</span> <span class="pl-c1">=</span> <span class="pl-s1">sequence_outputs</span>[<span class="pl-c1">1</span>]

        <span class="pl-k">return</span> <span class="pl-s1">lstm_out</span>

<span class="pl-c"># Instantiation:</span>
<span class="pl-s1">my_lstm</span> <span class="pl-c1">=</span> <span class="pl-v">MyLSTM</span>()</pre></div>
<p dir="auto">Fails with:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="for operator (Tensor 0, Tensor 1, Tensor 2, Tensor 3) -&gt; 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
                                                 ~~~~ &lt;--- HERE
                                                 True # enforce_sorted
                                                )"><pre class="notranslate"><code class="notranslate">for operator (Tensor 0, Tensor 1, Tensor 2, Tensor 3) -&gt; 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
                                                 ~~~~ &lt;--- HERE
                                                 True # enforce_sorted
                                                )
</code></pre></div>
<p dir="auto">So it seems that it expects Tensors for the arguments <code class="notranslate">batch_first</code> and <code class="notranslate">enforce_sorted</code> instead of Python bools. Okay, if we change them to:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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
                                            )"><pre class="notranslate"><span class="pl-s1">x_packed</span> <span class="pl-c1">=</span> <span class="pl-s1">nn</span>.<span class="pl-s1">utils</span>.<span class="pl-s1">rnn</span>.<span class="pl-en">pack_padded_sequence</span>(<span class="pl-s1">x_in_shortened</span>, 
                                             <span class="pl-s1">x_lengths</span>,
                                             <span class="pl-s1">torch</span>.<span class="pl-en">ones</span>(<span class="pl-c1">1</span>).<span class="pl-en">byte</span>().<span class="pl-en">squeeze</span>(), <span class="pl-c"># batch_first</span>
                                             <span class="pl-s1">torch</span>.<span class="pl-en">ones</span>(<span class="pl-c1">1</span>).<span class="pl-en">byte</span>().<span class="pl-en">squeeze</span>() <span class="pl-c"># enforce_sorted</span>
                                            )</pre></div>
<p dir="auto">Insantiation works, but it fails when the forward method is applied:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="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"><pre class="notranslate"><code class="notranslate">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
</code></pre></div>
<p dir="auto">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...</p>
<p dir="auto">BTW Thank you for PyTorch overall. ¡It's awesome!</p>