code: def testConstruct(self, constructor, kwargs):
      vqvae_module = constructor(**kwargs)
      # Batch of input vectors to quantize
      inputs_np = np.random.randn(100, kwargs["embedding_dim"]).astype(
          np.float32
      )
      inputs = jnp.array(inputs_np)
  
      # Set is_training to False, otherwise for the EMA case just evaluating the
      # forward pass will change the embeddings, meaning that some of our computed
      # closest embeddings will be incorrect.
      vq_output = vqvae_module(inputs, is_training=False)
  
      # Output shape is correct
      self.assertEqual(vq_output["quantize"].shape, inputs.shape)
  
      vq_output_np = jax.tree_util.tree_map(lambda t: t, vq_output)
      embeddings_np = vqvae_module.embeddings
  
      self.assertEqual(
          embeddings_np.shape,
          (kwargs["embedding_dim"], kwargs["num_embeddings"]),
      )
  
      # Check that each input was assigned to the embedding it is closest to.
      distances = (
          jnp.square(inputs_np).sum(axis=1, keepdims=True)
          - 2 * np.dot(inputs_np, embeddings_np)
          + jnp.square(embeddings_np).sum(axis=0, keepdims=True)
      )
      closest_index = np.argmax(-distances, axis=1)
      # On TPU, distances can be different by ~1% due to precision. This can cause
      # the distanc to the closest embedding to flip, leading to a difference
      # in the encoding indices tensor. First we check that the continuous
      # distances are reasonably close, and then we only allow N differences in
      # the encodings. For batch of 100, N == 3 seems okay (passed 1000x tests).
      np.testing.assert_allclose(
          distances, vq_output_np["distances"], atol=5e-2
      )
      num_differences_in_encodings = (
          closest_index != vq_output_np["encoding_indices"]
      ).sum()
      num_differences_allowed = 3
      self.assertLessEqual(
          num_differences_in_encodings, num_differences_allowed
      )
  
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 1 for calling the VqvaeTest obj's underlying member method code with example usage: obj.testConstruct(constructor, kwargs) without return types