def _compute_word_overlap(context_ids, context_len, question_ids, question_len,
                          reduce_type, weighted, vocab_df):
  """Compute word overlap between question and context ids.

  Args:
    context_ids: <int32> [batch_size, num_contexts, max_context_len]
    context_len: <int32> [batch_size, num_contexts]
    question_ids: <int32> [batch_size, max_question_len]
    question_len: <int32> [batch_size]
    reduce_type: String for reduce type when computing overlap. Choices are: max
      - Allows at most one match per question word. sum - Sums over all matches
      for each question word.
    weighted: Boolean indicate whether or not weight the overlap by IDF.
    vocab_df: Tensor of shape [vocab_size] for word frequency. Computes this at
      the document-level if not given.

  Returns:
    overlap: <float32> [batch_size, num_contexts]

  Raises:
    Exception: If invalid reduce_type is provided.
  """
  # <float> [batch_size, num_contexts, question_len, context_len]
  overlap = tf.to_float(
      _word_overlap_helper(question_ids=question_ids, context_ids=context_ids))

  # <float> [batch_size, question_len]
  question_mask = tf.sequence_mask(
      question_len, tf.shape(question_ids)[1], dtype=tf.float32)

  # <float> [batch_size, num_contexts, context_len]
  context_mask = tf.sequence_mask(
      context_len, tf.shape(context_ids)[2], dtype=tf.float32)

  overlap *= tf.expand_dims(tf.expand_dims(question_mask, 1), -1)
  overlap *= tf.expand_dims(context_mask, 2)

  if weighted:
    if vocab_df is None:
      # Use document-level IDF computed with respect to the current batch.
      flat_context_ids = tf.to_int32(tf.reshape(context_ids, [-1]))

      # <float> [number of unique words]
      vocab_df = tf.bincount(
          flat_context_ids,
          minlength=tf.reduce_max(question_ids) + 1,
          dtype=tf.float32)

      # Replace all zeros with ones.
      vocab_df = tf.where(
          tf.equal(vocab_df, 0), x=tf.ones_like(vocab_df), y=vocab_df)

    # <float>[batch_size, question_len] expanded to
    # <float> [batch_size, 1, question_len, 1]
    question_df = tf.gather(vocab_df, question_ids)
    question_df = tf.expand_dims(tf.expand_dims(question_df, 1), -1)

    # <float> [batch_size, num_contexts, question_len, context_len]
    overlap = tf.divide(tf.to_float(overlap), question_df)

  if reduce_type == "max":
    # <float> [batch_size, num_contexts]
    overlap = tf.reduce_sum(tf.reduce_max(overlap, axis=[3]), axis=[2])
  elif reduce_type == "sum":
    # <float> [batch_size, num_contexts]
    overlap = tf.reduce_sum(overlap, axis=[2, 3])
  else:
    raise Exception("Reduce type %s is invalid." % reduce_type)

  return overlap
Compute word overlap between question and context ids.

  Args:
    context_ids: <int32> [batch_size, num_contexts, max_context_len]
    context_len: <int32> [batch_size, num_contexts]
    question_ids: <int32> [batch_size, max_question_len]
    question_len: <int32> [batch_size]
    reduce_type: String for reduce type when computing overlap. Choices are: max
      - Allows at most one match per question word. sum - Sums over all matches
      for each question word.
    weighted: Boolean indicate whether or not weight the overlap by IDF.
    vocab_df: Tensor of shape [vocab_size] for word frequency. Computes this at
      the document-level if not given.

  Returns:
    overlap: <float32> [batch_size, num_contexts]

  Raises:
    Exception: If invalid reduce_type is provided.
  
language/question_answering/decatt_docreader/models/nq_long_decatt_model.py