def flatten_bert_inputs(bert_inputs):
  """Flatten all tensors in a BertInput and also return the inverse."""
  flat_token_ids, unflatten = tensor_utils.flatten(bert_inputs.token_ids)
  flat_mask, _ = tensor_utils.flatten(bert_inputs.mask)
  flat_segment_ids, _ = tensor_utils.flatten(bert_inputs.segment_ids)
  flat_bert_inputs = featurization.BertInputs(
      token_ids=flat_token_ids, mask=flat_mask, segment_ids=flat_segment_ids)
  return flat_bert_inputs, unflatten
Flatten all tensors in a BertInput and also return the inverse.
language/realm/model.py