def compute_loss(logits, indices, weights=None):
  """Cross-entropy loss when multiple labels per example."""
  if weights is not None:
    indices = indices * weights
    indices = indices / tf.maximum(
        1e-12, tf.reduce_sum(indices, axis=1, keepdims=True))
    row_sums = tf.reduce_sum(indices, axis=1)
  losses = tf.nn.softmax_cross_entropy_with_logits(
      labels=indices, logits=logits)
  if weights is not None:
    return tf.reduce_sum(row_sums * losses) / tf.maximum(
        tf.reduce_sum(row_sums), 1e-12), row_sums
  else:
    return tf.reduce_mean(losses), None
Cross-entropy loss when multiple labels per example.
language/labs/drkit/model_fns.py