Add support for fast variable-length LSTM

cpp-package/example/charRNN.cpp

@@ -164,8 +164,9 @@ Symbol LSTMWithBuiltInRNNOp(int num_lstm_layer, int sequence_length, int input_d
   auto rnn_h_init = Symbol::Variable("LSTM_init_h");
   auto rnn_c_init = Symbol::Variable("LSTM_init_c");
   auto rnn_params = Symbol::Variable("LSTM_parameters");  // See explanations near RNNXavier class
-  auto rnn = RNN(embed, rnn_params, rnn_h_init, rnn_c_init, num_hidden, num_lstm_layer,
-      RNNMode::kLstm, false, dropout, !isTrain);
+  auto variable_sequence_length = Symbol::Variable("sequence_length");
+  auto rnn = RNN(embed, rnn_params, rnn_h_init, rnn_c_init, variable_sequence_length, num_hidden,
+                 num_lstm_layer, RNNMode::kLstm, false, dropout, !isTrain);
   auto hidden = Reshape(rnn[0], Shape(), false, Shape(0, num_hidden), false);
 
   auto cls_weight = Symbol::Variable("cls_weight");

python/mxnet/gluon/rnn/rnn_layer.py

@@ -37,7 +37,7 @@ def __init__(self, hidden_size, num_layers, layout,
                  i2h_bias_initializer, h2h_bias_initializer,
                  mode, projection_size, h2r_weight_initializer,
                  lstm_state_clip_min, lstm_state_clip_max, lstm_state_clip_nan,
-                 dtype, **kwargs):
+                 dtype, use_sequence_length=False, **kwargs):
         super(_RNNLayer, self).__init__(**kwargs)
         assert layout in ('TNC', 'NTC'), \
             "Invalid layout %s; must be one of ['TNC' or 'NTC']"%layout
@@ -58,6 +58,7 @@ def __init__(self, hidden_size, num_layers, layout,
         self._lstm_state_clip_max = lstm_state_clip_max
         self._lstm_state_clip_nan = lstm_state_clip_nan
         self._dtype = dtype
+        self._use_sequence_length = use_sequence_length
 
         self._gates = {'rnn_relu': 1, 'rnn_tanh': 1, 'lstm': 4, 'gru': 3}[mode]
 
@@ -219,29 +220,39 @@ def begin_state(self, batch_size=0, func=ndarray.zeros, **kwargs):
             states.append(func(name='%sh0_%d'%(self.prefix, i), **info))
         return states
 
-    def hybrid_forward(self, F, inputs, states=None, **kwargs):
-        if F is ndarray:
-            batch_size = inputs.shape[self._layout.find('N')]
-        skip_states = states is None
-        if skip_states:
-            if F is ndarray:
+    def __call__(self, inputs, states=None, sequence_length=None, **kwargs):
+        self.skip_states = states is None
+        if states is None:
+            if isinstance(inputs, ndarray.NDArray):
+                batch_size = inputs.shape[self._layout.find('N')]
                 states = self.begin_state(batch_size, ctx=inputs.context, dtype=inputs.dtype)
             else:
                 states = self.begin_state(0, func=symbol.zeros)
         if isinstance(states, tensor_types):
             states = [states]
+
+        if self._use_sequence_length:
+            return super(_RNNLayer, self).__call__(inputs, states, sequence_length, **kwargs)
+        else:
+            return super(_RNNLayer, self).__call__(inputs, states, **kwargs)
+
+
+    def hybrid_forward(self, F, inputs, states, sequence_length=None, **kwargs):
+        if F is ndarray:
+            batch_size = inputs.shape[self._layout.find('N')]
+
         if F is ndarray:
             for state, info in zip(states, self.state_info(batch_size)):
                 if state.shape != info['shape']:
                     raise ValueError(
                         "Invalid recurrent state shape. Expecting %s, got %s."%(
                             str(info['shape']), str(state.shape)))
-        out = self._forward_kernel(F, inputs, states, **kwargs)
+        out = self._forward_kernel(F, inputs, states, sequence_length, **kwargs)
 
         # out is (output, state)
-        return out[0] if skip_states else out
+        return out[0] if self.skip_states else out
 
-    def _forward_kernel(self, F, inputs, states, **kwargs):
+    def _forward_kernel(self, F, inputs, states, sequence_length, **kwargs):
         """ forward using CUDNN or CPU kenrel"""
         if self._layout == 'NTC':
             inputs = F.swapaxes(inputs, dim1=0, dim2=1)
@@ -261,14 +272,20 @@ def _forward_kernel(self, F, inputs, states, **kwargs):
 
         params = F._internal._rnn_param_concat(*params, dim=0)
 
-        rnn = F.RNN(inputs, params, *states, state_size=self._hidden_size,
-                    projection_size=self._projection_size,
+        if self._use_sequence_length:
+            rnn_args = states + [sequence_length]
+        else:
+            rnn_args = states
+
+        rnn = F.RNN(inputs, params, *rnn_args, use_sequence_length=self._use_sequence_length,
+                    state_size=self._hidden_size, projection_size=self._projection_size,
                     num_layers=self._num_layers, bidirectional=self._dir == 2,
                     p=self._dropout, state_outputs=True, mode=self._mode,
                     lstm_state_clip_min=self._lstm_state_clip_min,
                     lstm_state_clip_max=self._lstm_state_clip_max,
                     lstm_state_clip_nan=self._lstm_state_clip_nan)
 
+
         if self._mode == 'lstm':
             outputs, states = rnn[0], [rnn[1], rnn[2]]
         else:

src/imperative/imperative_utils.h

@@ -67,6 +67,7 @@ inline Context GetContext(const nnvm::NodeAttrs& attrs,
   Context ctx;
   if (inputs.size()) {
     ctx = inputs[0]->ctx();
+
     for (size_t i = 1; i < inputs.size(); ++i) {
       CHECK_EQ(inputs[i]->ctx().dev_mask(), ctx.dev_mask())
           << "Operator " << attrs.op->name