def expand_example(features, sample_one=True):
  """Expand nested tensor protos into multiple examples."""
  question_ids = tf.io.parse_tensor(features["question_ids"], out_type=tf.int64)
  questions = tf.io.parse_tensor(features["questions"], out_type=tf.string)
  answers = tf.io.parse_tensor(features["answers"], out_type=tf.string)
  num_qas = tensor_utils.shape(questions, 0)

  if sample_one:
    rid = tf.random.uniform([], maxval=num_qas, dtype=tf.int32)
    question_ids = tf.expand_dims(question_ids[rid], 0)
    questions = tf.expand_dims(questions[rid], 0)
    answers = tf.expand_dims(answers[rid], 0)
    num_qas = 1

  image_ids = tf.tile(tf.expand_dims(features["image_id"], 0), [num_qas])
  images = tf.tile(tf.expand_dims(features["image"], 0), [num_qas])
  object_features = tf.tile(
      tf.expand_dims(features["object_features"], 0), [num_qas])
  object_positions = tf.tile(
      tf.expand_dims(features["object_positions"], 0), [num_qas])

  features = dict(
      image_id=image_ids,
      image=images,
      object_features=object_features,
      object_positions=object_positions,
      question_id=question_ids,
      question=questions,
      answer=answers)
  return tf.data.Dataset.from_tensor_slices(features)
Expand nested tensor protos into multiple examples.
language/capwap/datasets/vqa_dataset.py