def _num_context_mapper(keys_to_map, prefix="num_"):
  """Wrapper for mapper that computes number of contexts."""

  def _mapper(dataset):
    """Computes number of contexts."""
    for k in keys_to_map:
      size = tf.shape(dataset[k])[-1]
      dataset[prefix + k] = size
    return dataset

  return _mapper
Wrapper for mapper that computes number of contexts.
language/question_answering/decatt_docreader/models/nq_long_model.py