code: def testEmaUpdating(self, use_tf_function, dtype):
      if self.primary_device == "TPU" and dtype == tf.float64:
          self.skipTest("F64 not supported by TPU")
  
      embedding_dim = 6
      np_dtype = np.float64 if dtype is tf.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_tf_function:
          vqvae_module = tf.function(vqvae_module)
  
      batch_size = 16
  
      prev_embeddings = vqvae_module.embeddings.numpy()
  
      # Embeddings should change with every forwards pass if is_training == True.
      for _ in range(10):
          inputs = tf.random.normal(
              [batch_size, embedding_dim], dtype=dtype
          )
          vqvae_module(inputs, is_training=True)
          current_embeddings = vqvae_module.embeddings.numpy()
          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 = tf.random.normal(
              [batch_size, embedding_dim], dtype=dtype
          )
          vqvae_module(inputs, is_training=False)
          current_embeddings = vqvae_module.embeddings.numpy()
          self.assertTrue((current_embeddings == prev_embeddings).all())
  
file_path: sonnet/sonnet/src/nets/vqvae_test.py
imports: from absl.testing import parameterized
  import numpy
  from sonnet.src import test_utils
  from sonnet.src.nets import vqvae
  import tensorflow
  import tree
label_desc: Use image node_id 4 for calling the VqvaeTest obj's underlying member method code with example usage: obj.testEmaUpdating(use_tf_function, dtype) without return types