<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def softmax_cross_entropy_with_logits(logits, targets, vocab_dim, z_loss=0.0):
 """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. """

  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