code: def softmax_cross_entropy_with_logits(logits, targets, vocab_dim, z_loss=0.0):
    
    if logits.shape != targets.shape:
      raise ValueError(
          "logits shape must equal targets shape"
          "logits=%s targets=%s" % (logits.to_string, targets.to_string))
    if vocab_dim not in logits.shape.dims:
      raise ValueError("vocab_dim must be in logits.shape.dims")
    log_z = mtf.reduce_logsumexp(logits, vocab_dim)
    log_softmax = logits - log_z
    loss = mtf.negative(
        mtf.reduce_sum(log_softmax * targets, reduced_dim=vocab_dim))
    if z_loss != 0:
      loss += z_loss * mtf.square(log_z)
    return loss
docs: Per-example softmax loss.
  
    if z_loss is nonzero, we add a loss equal to z_loss*log(z)^2, where z is the
    partition function.  Example value: z_loss=1e-4.  Two uses of z_loss are:
    - To keep the logits from drifting too far from zero, which can cause
       unacceptable roundoff errors in bfloat16.
    - To encourage the logits to be normalized log-probabilities.
  
    Args:
      logits: a mtf.Tensor whose shape contains vocab_dim
      targets: a mtf.Tensor with the same shape as logits
      vocab_dim: a mtf.Dimension
      z_loss: a float
  
    Returns:
      a mtf.Tensor whose shape is equal to logits.shape - vocab_dim
  
    Raises:
      ValueError: if the shapes do not match.