code: def per_example_loss(labels, predictions, weights=None):
      """Calculate a per-example sparse categorical crossentropy loss.
  
      This loss function assumes that the predictions are post-softmax.
      Args:
        labels: The labels to evaluate against. Should be a set of integer indices
          ranging from 0 to (vocab_size-1).
        predictions: The network predictions. Should have softmax already applied.
        weights: An optional weight array of the same shape as the 'labels' array.
          If None, all examples will be used.
  
      Returns:
        A tensor of shape predictions.shape[:-1] containing the per-example
          loss.
      """
      # When using these functions with the Keras core API, we will need to squeeze
      # the labels tensor - Keras adds a spurious inner dimension.
      labels, predictions = _adjust_labels(labels, predictions)
      _validate_rank(labels, predictions, weights)
  
      labels_one_hot = tf.keras.backend.one_hot(
          labels, predictions.shape[-1]
      )
      labels_one_hot = tf.keras.backend.cast(
          labels_one_hot, predictions.dtype
      )
      per_example_loss_data = -tf.keras.backend.sum(
          predictions * labels_one_hot, axis=[-1]
      )
      if weights is not None:
          weights = tf.keras.backend.cast(
              weights, per_example_loss_data.dtype
          )
          per_example_loss_data = weights * per_example_loss_data
      return per_example_loss_data
  
file_path: autodist/examples/benchmark/utils/modeling/losses/weighted_sparse_categorical_crossentropy.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import tensorflow
label_desc: Use image node_id 3 for calling a global function with example usage: per_example_loss(labels, predictions, weights) and returns: per_example_loss_data