docs: def learning_rate_with_decay(
      batch_size, batch_denom, num_images, boundary_epochs, decay_rates,
      base_lr=0.1, enable_lars=False):
    """Get a learning rate that decays step-wise as training progresses.
  
    Args:
      batch_size: the number of examples processed in each training batch.
      batch_denom: this value will be used to scale the base learning rate.
        `0.1 * batch size` is divided by this number, such that when
        batch_denom == batch_size, the initial learning rate will be 0.1.
      num_images: total number of images that will be used for training.
      boundary_epochs: list of ints representing the epochs at which we
        decay the learning rate.
      decay_rates: list of floats representing the decay rates to be used
        for scaling the learning rate. It should have one more element
        than `boundary_epochs`, and all elements should have the same type.
      base_lr: Initial learning rate scaled based on batch_denom.
  
    Returns:
      Returns a function that takes a single argument - the number of batches
      trained so far (global_step)- and returns the learning rate to be used
      for training the next batch.
    """
    initial_learning_rate = base_lr * batch_size / batch_denom
    batches_per_epoch = num_images / batch_size
  
    # Multiply the learning rate by 0.1 at 100, 150, and 200 epochs.
    boundaries = [int(batches_per_epoch * epoch) for epoch in boundary_epochs]
    vals = [initial_learning_rate * decay for decay in decay_rates]
  
    def learning_rate_fn(global_step):
      lr = tf.train.piecewise_constant(global_step, boundaries, vals)
      warmup_steps = int(batches_per_epoch * 5)
      warmup_lr = (
          initial_learning_rate * tf.cast(global_step, tf.float32) / tf.cast(
          warmup_steps, tf.float32))
      return tf.cond(global_step < warmup_steps, lambda: warmup_lr, lambda: lr)
  
    def poly_rate_fn(global_step):
      """Handles linear scaling rule, gradual warmup, and LR decay.
  
      The learning rate starts at 0, then it increases linearly per step.  After
      flags.poly_warmup_epochs, we reach the base learning rate (scaled to account
      for batch size). The learning rate is then decayed using a polynomial rate
      decay schedule with power 2.0.
  
      Args:
      global_step: the current global_step
  
      Returns:
      returns the current learning rate
      """
  
      # Learning rate schedule for LARS polynomial schedule
      if batch_size < 8192:
        plr = 5.0
        w_epochs = 5
      elif batch_size < 16384:
        plr = 10.0
        w_epochs = 5
      elif batch_size < 32768:
        plr = 25.0
        w_epochs = 5
      else:
        plr = 32.0
        w_epochs = 14
  
      w_steps = int(w_epochs * batches_per_epoch)
      wrate = (plr * tf.cast(global_step, tf.float32) / tf.cast(
          w_steps, tf.float32))
  
      # TODO(pkanwar): use a flag to help calc num_epochs.
      num_epochs = 90
      train_steps = batches_per_epoch * num_epochs
  
      min_step = tf.constant(1, dtype=tf.int64)
      decay_steps = tf.maximum(min_step, tf.subtract(global_step, w_steps))
      poly_rate = tf.train.polynomial_decay(
          plr,
          decay_steps,
          train_steps - w_steps + 1,
          power=2.0)
      return tf.where(global_step <= w_steps, wrate, poly_rate)
  
    # For LARS we have a new learning rate schedule
    if enable_lars:
      return poly_rate_fn
  
    return learning_rate_fn