[jit] make nn.LSTM accept PackedSequence instead of Tuples

wanchaol · wanchaol · commit afe4b294da99 · 2019-07-31T19:57:39.000-07:00
Summary: Test Plan: Reviewers: Subscribers: Tasks: Tags: ghstack-source-id: 811724a Pull Request resolved: #23643
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -12786,6 +12786,7 @@ def forward(self, x, lengths, h0, c0):
         self.assertEqual(eager_out, script_out)
 
     def test_nn_LSTM(self):
+        from torch.nn.utils.rnn import PackedSequence
         input = torch.nn.utils.rnn.pack_sequence([torch.randn(5, 5)])
 
         class S(torch.jit.ScriptModule):
@@ -12795,7 +12796,7 @@ def __init__(self):
 
             @torch.jit.script_method
             def forward(self, input):
-                # type: (Tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor]]) -> Tuple[Tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor]], Tuple[Tensor, Tensor]]  # noqa
+                # type: (PackedSequence) -> Tuple[PackedSequence, Tuple[Tensor, Tensor]]  # noqa
                 return self.x(input)
 
         eager_out = self.runAndSaveRNG(lambda x: torch.nn.LSTM(5, 5)(x), (input,))[0]
diff --git a/torch/jit/__init__.py b/torch/jit/__init__.py
@@ -2002,7 +2002,6 @@ def register_all(mod):
         (torch.nn.init._no_grad_normal_, "aten::_no_grad_normal_"),
         (torch.nn.init._no_grad_uniform_, "aten::_no_grad_uniform_"),
         (torch.nn.init._no_grad_zero_, "aten::_no_grad_zero_"),
-        (torch.nn.utils.rnn.get_packed_sequence, "aten::_pack_sequence"),
         (torch._C._get_tracing_state, "aten::_get_tracing_state"),
         (warnings.warn, "aten::warn"),
     ]
diff --git a/torch/nn/modules/rnn.py b/torch/nn/modules/rnn.py
@@ -5,7 +5,7 @@
 
 from .module import Module
 from ..parameter import Parameter
-from ..utils.rnn import PackedSequence, get_packed_sequence
+from ..utils.rnn import PackedSequence
 from .. import init
 from .. import _VF
 from ..._jit_internal import _parameter_list
@@ -544,14 +544,14 @@ def forward_tensor(self, input, hx=None):
 
     @torch._jit_internal.export
     def forward_packed(self, input, hx=None):
-        # type: (Tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor]], Optional[Tuple[Tensor, Tensor]]) -> Tuple[Tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor]], Tuple[Tensor, Tensor]]  # noqa
+        # type: (PackedSequence, Optional[Tuple[Tensor, Tensor]]) -> Tuple[PackedSequence, Tuple[Tensor, Tensor]]  # noqa
         input, batch_sizes, sorted_indices, unsorted_indices = input
         max_batch_size = batch_sizes[0]
         max_batch_size = int(max_batch_size)
 
         output, hidden = self.forward_impl(input, hx, batch_sizes, max_batch_size, sorted_indices)
 
-        output = get_packed_sequence(output, batch_sizes, sorted_indices, unsorted_indices)
+        output = PackedSequence(output, batch_sizes, sorted_indices, unsorted_indices)
         return output, self.permute_hidden(hidden, unsorted_indices)
 
     @torch._jit_internal.ignore
diff --git a/torch/nn/utils/rnn.py b/torch/nn/utils/rnn.py
@@ -419,7 +419,3 @@ def pack_sequence(sequences, enforce_sorted=True):
     """
     lengths = [v.size(0) for v in sequences]
     return pack_padded_sequence(pad_sequence(sequences), lengths, enforce_sorted=enforce_sorted)
-
-
-def get_packed_sequence(data, batch_sizes, sorted_indices, unsorted_indices):
-    return PackedSequence(data, batch_sizes, sorted_indices, unsorted_indices)

Original file line number	Diff line number	Diff line change
`@@ -2002,7 +2002,6 @@ def register_all(mod):`
`2002`	`2002`	`(torch.nn.init._no_grad_normal_, "aten::_no_grad_normal_"),`
`2003`	`2003`	`(torch.nn.init._no_grad_uniform_, "aten::_no_grad_uniform_"),`
`2004`	`2004`	`(torch.nn.init._no_grad_zero_, "aten::_no_grad_zero_"),`
`2005`		`- (torch.nn.utils.rnn.get_packed_sequence, "aten::_pack_sequence"),`
`2006`	`2005`	`(torch._C._get_tracing_state, "aten::_get_tracing_state"),`
`2007`	`2006`	`(warnings.warn, "aten::warn"),`
`2008`	`2007`	`]`