text: """Training algorithm track submission functions for DLRM."""
  from typing import Iterator, List, Tuple, Union
  
  import torch
  import spec
  from . import workload
  
  
  def get_batch_size(workload_name:str):
      pass
  
  
  def init_optimizer_state(
      workload: spec.Workload,
      model_params: spec.ParameterContainer,
      model_state: spec.ModelAuxiliaryState,
      hyperparameters: spec.Hyperparamters,
      rng: spec.RandomState) -> spec.OptimizerState:
      pass
  
  def update_params(
      workload: spec.Workload,
      current_param_container: spec.ParameterContainer,
      current_params_types: spec.ParameterTypeTree,
      model_state: spec.ModelAuxiliaryState,
      hyperparameters: spec.Hyperparamters,
      input_batch: spec.Tensor,
      label_batch: spec.Tensor,
      # This will define the output activation via `output_activation_fn`.
      loss_type: spec.LossType,
      optimizer_state: spec.OptimizerState,
      eval_results: List[Tuple[int, float]],
      global_step: int,
      rng: spec.RandomState) -> spec.UpdateReturn:
      """Return (updated_optimizer_state, updated_params)."""
      pass
  
  # Not allowed to update the model parameters, hyperparameters, global step, or
  # optimzier state.
  def data_selection(
      workload: spec.Workload,
      input_queue: Iterator[Tuple[spec.Tensor, spec.Tensor]],
      optimizer_state: spec.OptimizerState,
      current_param_container: spec.ParameterContainer,
      hyperparameters: spec.Hyperparamters,
      global_step: int,
      rng: spec.RandomState) -> Tuple[spec.Tensor, spec.Tensor]:
      """Select data from the infinitely repeating, pre-shuffled input queue.
      Each element of the queue is a single training example and label.
      We left out `current_params_types` because we do not believe that it would
      # be necessary for this function.
      Return a tuple of input label batches.
      """
      pass
  