import tensorflow as tf

    return pt_loss


def log_sum_exp(x):
  """numerically stable log_sum_exp implementation that prevents overflow."""
  axis = len(x.get_shape()) - 1
  m = tf.reduce_max(x, axis)
  m2 = tf.reduce_max(x, axis, keep_dims=True)
  return m + tf.log(tf.reduce_sum(tf.exp(x - m2), axis))


def log_prob_from_logits(x):
  """numerically stable log_softmax implementation that prevents overflow."""
  axis = len(x.get_shape()) - 1
  m = tf.reduce_max(x, axis, keep_dims=True)
  return x - m - tf.log(tf.reduce_sum(tf.exp(x - m), axis, keep_dims=True))


def segment_loss(logits, labels, num_classes, head=None):
  """Calculate the loss from the logits and the labels.

  Args:
    logits: tensor, float - [batch_size * width * height, num_classes].
        Use vgg_fcn.up as logits.
    labels: Labels tensor, int32 - [batch_size * width * height, num_classes].
        The ground truth of your data.
    head: numpy array - [num_classes]
        Weighting the loss of each class
        Optional: Prioritize some classes

