def METHOD_NAME(optimizer_fn,
                      batch_per_epoch=1,
                      momentum=0.9,
                      renorm=False,
                      update_ops_in_replica_mode=False):
  """Example of non-distribution-aware legacy code with batch normalization."""
  def dataset_fn():
    # input shape is [16, 8], input values are increasing in both dimensions.
    return dataset_ops.Dataset.from_tensor_slices(
        [[[float(x * 8 + y + z * 100)
           for y in range(8)]
          for x in range(16)]
         for z in range(batch_per_epoch)]).repeat()
  optimizer = optimizer_fn()
  batchnorm = normalization.BatchNormalization(
      renorm=renorm, momentum=momentum, fused=False)
  layer = core.Dense(1, use_bias=False)
  def model_fn(x):
    """A model that uses batchnorm."""
    def loss_fn():
      y = batchnorm(x, training=True)
      with ops.control_dependencies(
          ops.get_collection(ops.GraphKeys.UPDATE_OPS)
          if update_ops_in_replica_mode else []):
        loss = math_ops.reduce_mean(
            math_ops.reduce_sum(layer(y)) - constant_op.constant(1.))
      # `x` and `y` will be fetched by the gradient computation, but not `loss`.
      return loss
    if isinstance(optimizer, optimizer_v2.OptimizerV2):
      return optimizer.minimize(loss_fn, lambda: layer.trainable_variables)
    # Callable loss.
    return optimizer.minimize(loss_fn)
  return model_fn, dataset_fn, batchnorm