aa1074dc1704d3732ab205c43d48083ef8c69680,tests/python/unittest/test_operator.py,,test_lstm_sym,#,120

Before Change

@with_seed()
@assert_raises_cudnn_not_satisfied(min_version="5.1.10")
def test_lstm_sym():
    T, N, I, H = 5, 32, 800, 800
    fused = mx.rnn.FusedRNNCell(H, num_layers=3, mode="lstm", get_next_state=True, prefix="")
    stack = mx.rnn.SequentialRNNCell()
    stack.add(mx.rnn.LSTMCell(H, prefix="l0_"))
    stack.add(mx.rnn.LSTMCell(H, prefix="l1_"))

After Change

@with_seed()
@assert_raises_cudnn_not_satisfied(min_version="5.1.10")
def test_lstm_sym():
    Ts = [1, 5]
    Ns = [1, 32]
    Is = [32, 128, 512]
    Hs = [32, 128, 512]
    for T, N, I, H in itertools.product(Ts, Ns, Is, Hs):
        fused = mx.rnn.FusedRNNCell(H, num_layers=3, mode="lstm", get_next_state=True, prefix="")
        stack = mx.rnn.SequentialRNNCell()
        stack.add(mx.rnn.LSTMCell(H, prefix="l0_"))
        stack.add(mx.rnn.LSTMCell(H, prefix="l1_"))
        stack.add(mx.rnn.LSTMCell(H, prefix="l2_"))

        check_rnn_consistency(fused, stack, T, N, I, H, "write")
        check_rnn_consistency(fused, stack, T, N, I, H, "add")
        check_rnn_consistency(fused, stack, T, N, I, H, "null")

@with_seed()
@assert_raises_cudnn_not_satisfied(min_version="5.1.10")
def test_lstm_bidirectional():
    Ts = [1, 5]

Instances