def truncate_contexts(context_token_ids, num_contexts, context_len,
                      max_contexts, max_context_len):
  """Truncate context based on max_contexts and max_context_words.

  Drops more than `max_contexts` contexts entirely, and truncates
  each remaining context to contain at most `max_context_words` words.

  Args:
    context_token_ids: <int32> [batch_size, padded_num_contexts,
      padded_context_len].
    num_contexts: <int32> [batch_size] for number of contexts in each batch.
    context_len: <int32> [batch_size, padded_num_contexts] for length of each
      context.
    max_contexts: <int32> for maximum number of contexts per example.
    max_context_len: <int32> for maximum length of each context.

  Returns:
    Truncated tensors for context_token_ids, context_len and document_len
    based on max_contexts and max_context_words.
  """
  # Compute effective max contexts and max context words.
  padded_num_contexts = tf.shape(context_token_ids)[1]
  padded_context_len = tf.shape(context_token_ids)[2]

  max_contexts = tf.to_int32(tf.minimum(padded_num_contexts, max_contexts))
  max_context_len = tf.to_int32(tf.minimum(padded_context_len, max_context_len))

  # Truncate context_token_ids, context_len, and document len
  context_token_ids = context_token_ids[:, :max_contexts, :max_context_len]
  context_len = context_len[:, :max_contexts]
  context_len = tf.minimum(context_len, max_context_len)
  num_contexts = tf.minimum(num_contexts, max_contexts)

  return context_token_ids, num_contexts, context_len
Truncate context based on max_contexts and max_context_words.

  Drops more than `max_contexts` contexts entirely, and truncates
  each remaining context to contain at most `max_context_words` words.

  Args:
    context_token_ids: <int32> [batch_size, padded_num_contexts,
      padded_context_len].
    num_contexts: <int32> [batch_size] for number of contexts in each batch.
    context_len: <int32> [batch_size, padded_num_contexts] for length of each
      context.
    max_contexts: <int32> for maximum number of contexts per example.
    max_context_len: <int32> for maximum length of each context.

  Returns:
    Truncated tensors for context_token_ids, context_len and document_len
    based on max_contexts and max_context_words.
  
language/question_answering/decatt_docreader/utils/nq_long_utils.py