def METHOD_NAME(self):
    batch_size, input_dim, output_dim = (4, 15, 15)
    x = np.ones((batch_size, 7, input_dim))
    layer = example.LinearBlock(output_dim)
    output = layer(x)
    self.assertAllEqual(output.shape, (batch_size, 7, output_dim))