def compute_loss_from_sptensors(sp_probs, sp_answers):
  """Cross-entropy loss when both answers and probabilities are sparse."""
  # First take log of the probabilities.
  mask = tf.greater(sp_probs.values, 1e-12)
  sp_probs = tf.sparse.retain(sp_probs, mask)
  sp_nll = tf.SparseTensor(
      indices=sp_probs.indices,
      values=-tf.log(sp_probs.values),
      dense_shape=sp_probs.dense_shape)
  return sp_sp_matmul(sp_nll, sp_answers)
Cross-entropy loss when both answers and probabilities are sparse.
language/labs/drkit/model_fns.py