def _multihead_attention_output_shape(  # pylint: disable=invalid-name
    input_shapes, feature_depth=None, **unused_kwargs):
  """Helper: calculate multihead attention output shape."""
  input_shape = input_shapes[0]  # Inputs are (q, k, v, mask).
  return input_shape[:-1] + (feature_depth,)