Function test_lstm

imperative/python/test/unit/module/test_rnn.py:129–172

(
    batch_size,
    seq_len,
    input_size,
    hidden_size,
    num_layers,
    bidirectional,
    init_hidden,
    batch_first,
)

    ],
)
def test_lstm(
    batch_size,
    seq_len,
    input_size,
    hidden_size,
    num_layers,
    bidirectional,
    init_hidden,
    batch_first,
):
    rnn = LSTM(
        input_size,
        hidden_size,
        batch_first=batch_first,
        num_layers=num_layers,
        bidirectional=bidirectional,
    )
    if batch_first:
        x_shape = (batch_size, seq_len, input_size)
    else:
        x_shape = (seq_len, batch_size, input_size)
    x = mge.random.normal(size=x_shape)
    total_hidden_size = num_layers * (2 if bidirectional else 1) * hidden_size
    if init_hidden:
        h = mge.random.normal(size=(batch_size, total_hidden_size))
        h = (h, h)
    else:
        h = None
    output, h_n = rnn(x, h)
    num_directions = 2 if bidirectional else 1
    if batch_first:
        assert_tuple_equal(
            output.shape, (batch_size, seq_len, num_directions * hidden_size)
        )
    else:
        assert_tuple_equal(
            output.shape, (seq_len, batch_size, num_directions * hidden_size)
        )
    assert_tuple_equal(
        h_n[0].shape, (num_directions * num_layers, batch_size, hidden_size)
    )
    assert_tuple_equal(
        h_n[1].shape, (num_directions * num_layers, batch_size, hidden_size)
    )