import tensorflow as tf

    original_shape: Shape of `labels` and `logits` before reshape.
  Raises:
    ValueError: If `labels` and `logits` do not have the same shape.
  """
  # Convert `labels` and `logits` to Tensors and standardize dtypes.
  logits = tf.convert_to_tensor(logits, name='logits')
  labels = losses_utils.convert_and_cast(labels, 'labels', logits.dtype.base_dtype)
  weights = losses_utils.convert_and_cast(weights, 'weights', logits.dtype.base_dtype)

  try:
