@parameterized.parameters([((3,),), (3,)])
def test_conv_transpose(self, kernel_size):
    rng = dict(params=random.PRNGKey(0))
    x = jnp.ones((1, 8, 3))
    conv_transpose_module = nn.ConvTranspose(features=4, kernel_size=kernel_size, padding='VALID', kernel_init=initializers.ones, bias_init=initializers.ones)
    (y, initial_params) = conv_transpose_module.init_with_output(rng, x)
    self.assertEqual(initial_params['params']['kernel'].shape, (3, 3, 4))
    correct_ans = np.array([[[4.0, 4.0, 4.0, 4.0], [7.0, 7.0, 7.0, 7.0], [10.0, 10.0, 10.0, 10.0], [10.0, 10.0, 10.0, 10.0], [10.0, 10.0, 10.0, 10.0], [10.0, 10.0, 10.0, 10.0], [10.0, 10.0, 10.0, 10.0], [10.0, 10.0, 10.0, 10.0], [7.0, 7.0, 7.0, 7.0], [4.0, 4.0, 4.0, 4.0]]])
    np.testing.assert_allclose(y, correct_ans)
@parameterized.parameters([((3,),), (3,)])
def test_conv_transpose(self, kernel_size):
    rng = dict(params=random.PRNGKey(0))
    x = jnp.ones((1, 8, 3))
    conv_transpose_module = nn.ConvTranspose(features=4, kernel_size=kernel_size, padding='VALID', kernel_init=initializers.ones, bias_init=initializers.ones)
    (y, initial_params) = conv_transpose_module.init_with_output(rng, x)
    self.assertEqual(initial_params['params']['kernel'].shape, (3, 3, 4))
    correct_ans = np.array([[[4.0, 4.0, 4.0, 4.0], [7.0, 7.0, 7.0, 7.0], [10.0, 10.0, 10.0, 10.0], [10.0, 10.0, 10.0, 10.0], [10.0, 10.0, 10.0, 10.0], [10.0, 10.0, 10.0, 10.0], [10.0, 10.0, 10.0, 10.0], [10.0, 10.0, 10.0, 10.0], [7.0, 7.0, 7.0, 7.0], [4.0, 4.0, 4.0, 4.0]]])
    np.testing.assert_allclose(y, correct_ans)