code: def loss(params, batch, model_predict, rng):
    
    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)
docs: Calculate loss.