code: def test_different_input_shapes(self):
      strategy, mid_level_api, _ = self._create_strategy_and_mid_level(
          "sgd"
      )
  
      sparse = self._create_high_dimensional_sparse_dataset(strategy)
      sparse_iter = iter(
          strategy.experimental_distribute_dataset(
              sparse,
              options=distribute_lib.InputOptions(
                  experimental_fetch_to_device=False
              ),
          )
      )
      # Create a feature with shape (1, 3, 1)
      dense_feature = constant_op.constant(
          np.zeros(3), shape=(1, 3, 1), dtype=dtypes.int32
      )
      dense_dataset = (
          dataset_ops.DatasetV2.from_tensors(dense_feature)
          .unbatch()
          .repeat()
          .batch(1 * strategy.num_replicas_in_sync, drop_remainder=True)
      )
      dense_iter = iter(
          strategy.experimental_distribute_dataset(
              dense_dataset,
              options=distribute_lib.InputOptions(
                  experimental_fetch_to_device=False
              ),
          )
      )
  
      @def_function.function
      def test_fn():
          def step():
              return mid_level_api.dequeue()
  
          features = (
              next(dense_iter),
              next(sparse_iter)[1],
              next(sparse_iter)[2],
          )
          mid_level_api.enqueue(features, training=False)
          return strategy.run(step)
  
      test_fn()
  
      self.assertEqual(
          mid_level_api._output_shapes,
          [
              TensorShape((1, 3)),
              TensorShape((self.batch_size, self.data_batch_size)),
              TensorShape((self.batch_size, self.data_batch_size)),
          ],
      )
  
file_path: tensorflow-upstream/tensorflow/python/tpu/tests/tpu_embedding_v2_hd_valid_input_test.py
imports: 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 constant_op
  from tensorflow.python.framework import dtypes
  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_different_input_shapes() and returns: strategy, mid_level_api