f METHOD_NAME(
  self,
  params: spec.ParameterContainer,
  augmented_and_preprocessed_input_batch: Dict[str, spec.Tensor],
  model_state: spec.ModelAuxiliaryState,
  mode: spec.ForwardPassMode,
  rng: spec.RandomState,
  update_batch_norm: bool) -> Tuple[spec.Tensor, spec.ModelAuxiliaryState]: