code: def test_pass_none_to_apply_gradients(self):
      self.skip_if_oss()
      strategy, mid_level_api, _ = self._create_strategy_and_mid_level(
          "sgd"
      )
      mid_level_api.build(
          [
              TensorShape((self.batch_size, 2)),
              TensorShape((self.batch_size, 2)),
              TensorShape((self.batch_size, 3)),
          ]
      )
      dataset = self._create_sparse_dataset(strategy)
      data = next(
          iter(
              strategy.experimental_distribute_dataset(
                  dataset,
                  options=distribute_lib.InputOptions(
                      experimental_fetch_to_device=False
                  ),
              )
          )
      )
  
      @def_function.function
      def embedding_and_set_gradients(data):
          mid_level_api.enqueue(data)
  
          def tpu_fn():
              results = mid_level_api.dequeue()
              mid_level_api.apply_gradients(
                  (None, None, array_ops.ones_like(results[2]))
              )
              return results
  
          return strategy.run(tpu_fn)
  
      @def_function.function
      def embedding_only(data):
          mid_level_api.enqueue(data, training=False)
  
          def tpu_fn():
              return mid_level_api.dequeue()
  
          return strategy.run(tpu_fn)
  
      first = self._get_replica_numpy(
          embedding_and_set_gradients(data), strategy, 0
      )
      second = self._get_replica_numpy(
          embedding_only(data), strategy, 0
      )
  
      # First two features should be the same as None gradient was applied.
      # Third feature had gradient of 1 passed in from each core.
      # Each core received the same ids per core and returned the following batch:
      # [ row 3, row 0 + row 1 + row 2 ]
      # so gradient update was (learning rate = 0.1):
      #   row 0: -1/3*0.1
      #   row 1: -1/3*0.1
      #   row 2: -1/3*0.1
      #   row 3: -1*0.1
      # There is a factor of num_replicas because each replica gave an update.
  
      num_replicas = strategy.num_replicas_in_sync
      update = (
          [[0.0]],
          [[0.0]],
          [[0.1 * num_replicas], [0.1 / 3 * num_replicas]],
      )
      golden = tuple(
          [feature - np.array(up) for feature, up in zip(first, update)]
      )
  
      self.assertAllClose(golden, second)
  
imports: from absl.testing import parameterized
  import numpy
  from tensorflow.python.compat import v2_compat
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.distribute import distribute_lib
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import sparse_tensor
  from tensorflow.python.framework.tensor_shape import TensorShape
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import init_ops_v2
  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.tpu.tests import tpu_embedding_base_test
  from tensorflow.python.util import nest
label_desc: Use image node_id 1 for calling the TPUEmbeddingTest obj's underlying member method code with example usage: obj.test_pass_none_to_apply_gradients() and returns: strategy, strategy, results, mid_level_api