code: def test_sequence_feature_with_build(self, is_updated_shape):
      seq_length = 3
      # Set the max_seq_length in feature config
      for feature in self.feature_config:
          feature.max_sequence_length = seq_length
      strategy, mid_level_api, _ = self._create_strategy_and_mid_level(
          "sgd"
      )
      dataset = self._create_sparse_dataset(strategy)
      feature_iter = iter(
          strategy.experimental_distribute_dataset(
              dataset,
              options=distribute_lib.InputOptions(
                  experimental_fetch_to_device=False
              ),
          )
      )
      if is_updated_shape:
          mid_level_api.build(
              [
                  TensorShape([self.batch_size, seq_length, 2]),
                  TensorShape([self.batch_size, seq_length, 2]),
                  TensorShape([self.batch_size, seq_length, 3]),
              ]
          )
      else:
          mid_level_api.build(
              [
                  TensorShape([self.batch_size, 2]),
                  TensorShape([self.batch_size, 2]),
                  TensorShape([self.batch_size, 3]),
              ]
          )
  
      @def_function.function
      def test_fn():
          def step():
              return mid_level_api.dequeue()
  
          mid_level_api.enqueue(next(feature_iter), training=False)
          return strategy.run(step)
  
      output = test_fn()
      self.assertEqual(
          self._get_replica_numpy(output[0], strategy, 0).shape,
          (2, 3, 4),
      )
      self.assertEqual(
          self._get_replica_numpy(output[1], strategy, 0).shape,
          (2, 3, 4),
      )
      self.assertEqual(
          self._get_replica_numpy(output[2], strategy, 0).shape,
          (2, 3, 2),
      )
  
file_path: tensorflow/tensorflow/python/tpu/tests/tpu_embedding_v2_sequence_feature_test.py
imports: from absl.testing import parameterized
  from tensorflow.python.compat import v2_compat
  from tensorflow.python.distribute import distribute_lib
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework.tensor_shape import TensorShape
  from tensorflow.python.platform import test
  from tensorflow.python.tpu.tests import tpu_embedding_base_test
label_desc: Use image node_id 2 for calling the TPUEmbeddingTest obj's underlying member method code with example usage: obj.test_sequence_feature_with_build(is_updated_shape) and returns: strategy, mid_level_api