f METHOD_NAME(self,
                       split: str,
                       num_examples: int,
                       global_batch_size: int,
                       params: spec.ParameterContainer,
                       model_state: spec.ModelAuxiliaryState,
                       rng: spec.RandomState,
                       data_dir: str,
                       global_step: int = 0) -> Dict[str, float]: