code: def _softmax_cross_entropy_with_logits(  # pylint: disable=invalid-name,unused-argument
      labels, logits, axis=-1, name=None
  ):
      """Softmax cross entropy with logits."""
      cost = -np.sum(
          np.where(
              labels == 0,
              np.zeros_like(labels),
              labels
              * (
                  logits
                  - reduce_logsumexp(logits, axis=axis, keepdims=True)
              ),
          ),
          axis=axis,
      )
      return cost.astype(logits.dtype)
  
file_path: probability/tensorflow_probability/python/internal/backend/numpy/nn.py
imports: import collections
  import numpy
  from tensorflow_probability.python.internal.backend.numpy import _utils
  from tensorflow_probability.python.internal.backend.numpy import numpy_array
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import l2_normalize
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import log_softmax
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import reduce_logsumexp
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import reduce_mean
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import softmax
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import softplus
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import squared_difference
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import top_k
  from tensorflow_probability.python.internal.backend.numpy.ops import stop_gradient
label_desc: Use image node_id 5 for calling a global function with example usage: _softmax_cross_entropy_with_logits(labels, logits, axis, name) and returns: cost