Function test_rnn_cell

imperative/python/test/unit/module/test_rnn.py:22–30

(batch_size, input_size, hidden_size, init_hidden)

    [(3, 10, 20, True), (3, 10, 20, False), (1, 10, 20, False)],
)
def test_rnn_cell(batch_size, input_size, hidden_size, init_hidden):
    rnn_cell = RNNCell(input_size, hidden_size)
    x = mge.random.normal(size=(batch_size, input_size))
    if init_hidden:
        h = F.zeros(shape=(batch_size, hidden_size))
    else:
        h = None
    h_new = rnn_cell(x, h)
    assert_tuple_equal(h_new.shape, (batch_size, hidden_size))


@pytest.mark.skipif(get_device_count("gpu") > 0, reason="no algorithm on cuda")