def test_connect_conv_same(self, n):
batch_size = 2
input_shape = (16,) * n
input_shape_b = (batch_size,) + input_shape + (4,)
data = jnp.zeros(input_shape_b)
core = recurrent.ConvNDLSTM(
n, input_shape=input_shape, output_channels=3, kernel_shape=3)
state = core.initial_state(batch_size=batch_size)
out, state = core(data, state)
expected_output_shape = (batch_size,) + input_shape + (3,)
self.assertEqual(out.shape, expected_output_shape)
self.assertEqual(state[0].shape, expected_output_shape)
self.assertEqual(state[1].shape, expected_output_shape)
def f(data, state_init):
net = recurrent.ConvNDLSTM(n,
input_shape=input_shape,
output_channels=3,
kernel_shape=3)
return net(data, state_init)
