code: def testShapeChecking(self, constructor, kwargs):
      vqvae_module = constructor(**kwargs)
      wrong_shape_input = np.random.randn(
          100, kwargs["embedding_dim"] * 2
      )
      with self.assertRaises((ValueError, TypeError)):
          vqvae_module(
              jnp.array(wrong_shape_input.astype(np.float32)),
              is_training=False,
          )
  
file_path: dm-haiku/haiku/_src/nets/vqvae_test.py
imports: import functools
  from absl.testing import absltest
  from absl.testing import parameterized
  from haiku._src import stateful
  from haiku._src import test_utils
  from haiku._src import transform
  from haiku._src.nets import vqvae
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 2 for calling the VqvaeTest obj's underlying member method code with example usage: obj.testShapeChecking(constructor, kwargs) without return types