f METHOD_NAME(
  self,
  params: spec.ParameterContainer,
  batch: Dict[str, spec.Tensor],
  model_state: spec.ModelAuxiliaryState,
  rng: spec.RandomState) -> Dict[spec.Tensor, spec.ModelAuxiliaryState]: