<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def loss(params, batch, model_predict, rng):
 """Calculate loss."""

  inputs, targets = batch
  predictions = model_predict(inputs, params, rng=rng)
  predictions, targets = _make_list(predictions, targets)
  xent = []
  for (pred, target) in zip(predictions, targets):
    xent.append(np.sum(pred * layers.one_hot(target, pred.shape[-1]), axis=-1))
  return - masked_mean(xent, targets)