code: def test_enqueue_per_device(self):
      self.skip_if_oss()
      strategy, mid_level_api, _ = self._create_strategy_and_mid_level(
          "sgd"
      )
  
      sparse = self._create_sparse_dataset(strategy)
      sparse_iter = iter(
          strategy.experimental_distribute_dataset(
              sparse,
              options=distribute_lib.InputOptions(
                  experimental_fetch_to_device=False
              ),
          )
      )
  
      @def_function.function
      def test_fn():
          def get_activations(dense_value):
              return mid_level_api.dequeue(), dense_value
  
          sparse_features = next(sparse_iter)
          mid_level_api.enqueue(sparse_features, training=False)
          activations, dense_value1 = strategy.run(
              get_activations, args=(0.0,)
          )
  
          def enqueue_fn(ctx):
              core_id = ctx.replica_id_in_sync_group
              device = strategy.extended.worker_devices[core_id]
              sparse_features_local = nest.map_structure(
                  lambda x: strategy.experimental_local_results(x)[
                      core_id
                  ],
                  sparse_features,
              )
              mid_level_api.enqueue(
                  sparse_features_local, training=False, device=device
              )
              return 0.0
  
          data = strategy.experimental_distribute_values_from_function(
              enqueue_fn
          )
          per_device_activations, dense_value2 = strategy.run(
              get_activations, args=(data,)
          )
          return (
              activations,
              per_device_activations,
              dense_value1,
              dense_value2,
          )
  
      activations, per_device_activations, _, _ = test_fn()
  
      # Extact per core numpy arrays and check that both sparse and ragged have
      # the same results.
      activations0 = self._get_replica_numpy(activations, strategy, 0)
      per_device_activations0 = self._get_replica_numpy(
          per_device_activations, strategy, 0
      )
      self.assertAllClose(activations0, per_device_activations0)
      test_fn()
  
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 3 for calling the TPUEmbeddingTest obj's underlying member method code with example usage: obj.test_enqueue_per_device() and returns: activations,  per_device_activations,  dense_value1,  dense_value2, mid_level_api,  dense_value, int