(batch_size, input_size, hidden_size, init_hidden)
| 36 | [(3, 10, 20, True), (3, 10, 20, False), (1, 10, 20, False)], |
| 37 | ) |
| 38 | def test_lstm_cell(batch_size, input_size, hidden_size, init_hidden): |
| 39 | rnn_cell = LSTMCell(input_size, hidden_size) |
| 40 | x = mge.random.normal(size=(batch_size, input_size)) |
| 41 | if init_hidden: |
| 42 | h = F.zeros(shape=(batch_size, hidden_size)) |
| 43 | hx = (h, h) |
| 44 | else: |
| 45 | hx = None |
| 46 | h_new, c_new = rnn_cell(x, hx) |
| 47 | assert_tuple_equal(h_new.shape, (batch_size, hidden_size)) |
| 48 | assert_tuple_equal(c_new.shape, (batch_size, hidden_size)) |
| 49 | |
| 50 | |
| 51 | @pytest.mark.skipif(get_device_count("gpu") > 0, reason="no algorithm on cuda") |
nothing calls this directly
no test coverage detected