code: def _check_embedding_and_slot_variables_for_adagrad_momentum(
      self, embedding_table_before, gradients, optimizer, variable
  ):
      embedding_table = np.copy(embedding_table_before)
      accumulator = np.zeros(
          self._get_variable(variable["accumulators"]).shape
      )
      momenta = np.zeros(self._get_variable(variable["momenta"]).shape)
      gradients = np.sum(gradients, axis=0)
      if optimizer.beta2 == 1.0:
          accumulator += gradients**2
      else:
          accumulator = (
              optimizer.beta2 * accumulator
              + (1 - optimizer.beta2) * gradients**2
          )
      accumulator_power = np.power(
          accumulator + optimizer.epsilon, -1.0 / optimizer.exponent
      )
      momenta = (
          optimizer.momentum * momenta + gradients * accumulator_power
      )
      if optimizer.use_nesterov:
          update = (
              optimizer.momentum * momenta
              + gradients * accumulator_power
          )
      else:
          update = momenta
      embedding_table -= optimizer.learning_rate * update
      self.assertAllClose(
          self._get_variable(variable["parameters"]).numpy(),
          embedding_table,
          rtol=1e-3,
      )
      self.assertAllClose(
          self._get_variable(variable["accumulators"]).numpy(),
          accumulator,
          rtol=1e-3,
      )
      self.assertAllClose(
          self._get_variable(variable["momenta"]).numpy(),
          momenta,
          rtol=1e-3,
      )
  
imports: import os
  from typing import Tuple
  from absl import flags
  from absl.testing import parameterized
  import numpy
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.distribute import tpu_strategy
  from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
  from tensorflow.python.eager import remote
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import sparse_tensor
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import gen_math_ops
  from tensorflow.python.ops import init_ops_v2
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops.ragged import ragged_tensor
  from tensorflow.python.platform import test
  from tensorflow.python.tpu import tpu_embedding_v2
  from tensorflow.python.tpu import tpu_embedding_v2_utils
  from tensorflow.python.util import nest
label_desc: Use image node_id 18 for calling the TPUEmbeddingBaseTest obj's underlying member method code with example usage: obj._check_embedding_and_slot_variables_for_adagrad_momentum(embedding_table_before, gradients, optimizer, variable) without return types