code: def testEmaUpdating(self, use_jit, dtype):
      if (
          jax.local_devices()[0].platform == "tpu"
          and dtype == jnp.float64
      ):
          self.skipTest("F64 not supported by TPU")
  
      embedding_dim = 6
      np_dtype = np.float64 if dtype is jnp.float64 else np.float32
      decay = np.array(0.1, dtype=np_dtype)
      vqvae_module = vqvae.VectorQuantizerEMA(
          embedding_dim=embedding_dim,
          num_embeddings=7,
          commitment_cost=0.5,
          decay=decay,
          dtype=dtype,
      )
  
      if use_jit:
          vqvae_f = stateful.jit(vqvae_module, static_argnums=1)
      else:
          vqvae_f = vqvae_module
  
      batch_size = 16
  
      prev_embeddings = vqvae_module.embeddings
  
      # Embeddings should change with every forwards pass if is_training == True.
      for _ in range(10):
          inputs = np.random.rand(batch_size, embedding_dim).astype(
              dtype
          )
          vqvae_f(inputs, True)
          current_embeddings = vqvae_module.embeddings
          self.assertFalse(
              (prev_embeddings == current_embeddings).all()
          )
          prev_embeddings = current_embeddings
  
      # Forward passes with is_training == False don't change anything
      for _ in range(10):
          inputs = np.random.rand(batch_size, embedding_dim).astype(
              dtype
          )
          vqvae_f(inputs, False)
          current_embeddings = vqvae_module.embeddings
          self.assertTrue((current_embeddings == prev_embeddings).all())
  
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 4 for calling the VqvaeTest obj's underlying member method code with example usage: obj.testEmaUpdating(use_jit, dtype) without return types