| 1898 | assert_dtype(layer, mx.int16) |
| 1899 | |
| 1900 | def test_rnn(self): |
| 1901 | layer = nn.RNN(input_size=5, hidden_size=12, bias=True) |
| 1902 | inp = mx.random.normal((2, 25, 5)) |
| 1903 | |
| 1904 | h_out = layer(inp) |
| 1905 | self.assertEqual(h_out.shape, (2, 25, 12)) |
| 1906 | |
| 1907 | layer = nn.RNN( |
| 1908 | 5, |
| 1909 | 12, |
| 1910 | bias=False, |
| 1911 | nonlinearity=lambda x: mx.maximum(0, x), |
| 1912 | ) |
| 1913 | |
| 1914 | h_out = layer(inp) |
| 1915 | self.assertEqual(h_out.shape, (2, 25, 12)) |
| 1916 | |
| 1917 | with self.assertRaises(ValueError): |
| 1918 | nn.RNN(5, 12, nonlinearity="tanh") |
| 1919 | |
| 1920 | inp = mx.random.normal((44, 5)) |
| 1921 | h_out = layer(inp) |
| 1922 | self.assertEqual(h_out.shape, (44, 12)) |
| 1923 | |
| 1924 | h_out = layer(inp, hidden=h_out[-1, :]) |
| 1925 | self.assertEqual(h_out.shape, (44, 12)) |
| 1926 | |
| 1927 | def test_gru(self): |
| 1928 | layer = nn.GRU(5, 12, bias=True) |