Bug in BidirectionalCell for imperative mode when time_step is 1

[karan6181]

The _reverse_sequences() method in ~/mxnet/gluon/rnn/rnn_cell.py gives unexpected output when the input is of single time_step with the imperative mode. Below is the example to demonstrate the issue:

import mxnet as mx
from mxnet import gluon

def test_bidirectional_unroll_valid_length():
    class BiLSTM(gluon.nn.HybridBlock):
        def __init__(self, rnn_size, time_step, **kwargs):
            super(BiLSTM, self).__init__(**kwargs)
            self.time_step = time_step
            with self.name_scope():
                self.bi_lstm = gluon.rnn.BidirectionalCell(
                    gluon.rnn.LSTMCell(rnn_size, prefix='rnn_l0_'),
                    gluon.rnn.LSTMCell(rnn_size, prefix='rnn_r0_'),
                    output_prefix='lstm_bi_')

        def hybrid_forward(self, F, inputs, valid_len):
            outputs, states = self.bi_lstm.unroll(self.time_step, inputs, valid_length=valid_len, layout='NTC', merge_outputs=True)
            return outputs, states

    rnn_size, time_step = 100, 1
    net = BiLSTM(rnn_size, time_step)
    net.initialize()
    # net.hybridize()
    inputs_data = mx.nd.random.uniform(shape=(10, 1, 50))
    valid_len = mx.nd.array([1]*10)
    outputs, _ = net(inputs_data, valid_len)
    assert outputs.shape == (10, 1, 200)


test_bidirectional_unroll_valid_length()

Below is the debugged output from _reverse_sequences() method

>>> len(sequences) # Before reversing the input
1
>>> sequences[0].shape # Before reversing the input
(10, 50)
>>> reversed_sequences.shape # After reversing the input
(10, 50)
>>> len(reversed_sequences) # After reversing the input
10

# The correct output should be
>>> len(reversed_sequences)
1
reversed_sequences[0].shape
(10, 50)

The issue is not reproducible in symbolic mode(Please uncomment the net.hybridize() if you want to run it in symbolic mode). Also, the failure can only be seen in the nightly pip release due to the change made in this PR #13575.

[mxnet-label-bot]

Hey, this is the MXNet Label Bot.
Thank you for submitting the issue! I will try and suggest some labels so that the appropriate MXNet community members can help resolve it.
Here are my recommended labels: Bug

[eric-haibin-lin]

@szha

[roywei]

@mxnet-label-bot add[Bug, Gluon]

[szha]

The problem is in the inconsistent behavior of split, where if the number of splits is 1 it returns a single NDArray instead of a list.