def METHOD_NAME(  # pylint: disable=invalid-name
    batch_size=2,
    weight_init_value=0.8,
    optimizer=None,
    learning_rate=1.0,
    max_train_steps=10,
    train_steps_per_loop=2,
    eval_decode_steps_per_loop=2,
    eval_decode_samples_per_summary=10):
  """Register an IdentityRegressionTask model with given configuration.
  Args:
    batch_size: batch size of CountingInputGenerator.
    weight_init_value: constant init value for the model varialbes.
    optimizer: if set, the optimizer params to use.
    learning_rate: the learning rate to use.
    max_train_steps: maximum training steps.
    train_steps_per_loop: number of training steps per TPU loop.
    eval_decode_steps_per_loop: number of evaluation/decode steps per TPU loop.
    eval_decode_samples_per_summary: number of samples to eval/decode for each
      checkpoint.
  """
  class IdentityRegressionModel(base_model_params.SingleTaskModelParams):
    """Model params for IdentityRegressionTask."""
    def Train(self):
      return CountingInputGenerator.Params().Set(batch_size=batch_size)
    def Test(self):
      return CountingInputGenerator.Params().Set(batch_size=batch_size)
    def Task(self):
      p = IdentityRegressionTask.Params().Set(
          weight_init_value=weight_init_value)
      if optimizer:
        p.train.optimizer = optimizer
      p.train.learning_rate = learning_rate
      p.train.max_steps = max_train_steps
      p.train.tpu_steps_per_loop = train_steps_per_loop
      p.eval.samples_per_summary = eval_decode_samples_per_summary
      p.eval.decoder_samples_per_summary = eval_decode_samples_per_summary
      return p
    def Model(self):
      return ModelTrackingFPropResults.Params(self.Task())
    def ProgramSchedule(self):
      p = program.SimpleProgramScheduleForTask(
          train_dataset_name='Train',
          train_steps_per_loop=train_steps_per_loop,
          eval_dataset_names=['Test'],
          eval_steps_per_loop=eval_decode_steps_per_loop,
          decode_steps_per_loop=eval_decode_steps_per_loop)
      if max_train_steps == 0:
        p.train_executions_per_eval = 0
      return p
  model_registry.RegisterSingleTaskModel(IdentityRegressionModel)