f METHOD_NAME(
  self,
  label_batch: spec.Tensor,  # Dense or one-hot labels.
  logits_batch: spec.Tensor,
  mask_batch: Optional[spec.Tensor] = None,
  label_smoothing: float = 0.0) -> Dict[str, spec.Tensor]:  # differentiable