(self)
| 12 | class TestFlax(unittest.TestCase): |
| 13 | |
| 14 | def test_pooling(self): |
| 15 | x = jnp.full((1, 3, 3, 1), 2.) |
| 16 | mul_reduce = lambda x, y: x * y |
| 17 | y = nn.pooling.pool(x, 1., mul_reduce, (2, 2), (1, 1), 'VALID') |
| 18 | np.testing.assert_allclose(y, np.full((1, 2, 2, 1), 2. ** 4)) |
| 19 | |
| 20 | def test_cnn(self): |
| 21 | class CNN(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected