class_docstr: An  Adam optimizer that includes "correct" L2 weight decay.
  
  Adam optimizer that is able to processes multiple learning rate schedules
  for different variables within the optimizer class itself. This function
  takes in a list of variables to search and a list of corresponding
  alt lr schedules.
  
  The original variables are processed with the original learning rate
  controlled from outside the loop.
  
  Learning rate schedule should use the product learning rate.
code: def apply_grad(self, grad, var):
      if self.alt_lr_schedules is None or self.variable_search is None:
          return super(AdamWithMultiLRSchedule, self).apply_grad(
              grad, var
          )
  
      actual_lr_rates = compute_lr_for_step(
          self.alt_lr_schedules, self.learning_rate
      )
  
      # Modify learning rate for exception variables
      for idx, variable_search in enumerate(self.variable_search):
          if re.search(variable_search, var.name) is not None:
              # finds variable in LR schedule
              old_lr = self.learning_rate
              # get n-th learning rate schedule
              self.learning_rate = actual_lr_rates[idx]
              assignments = super(
                  AdamWithMultiLRSchedule, self
              ).apply_grad(grad, var)
              self.learning_rate = old_lr
          else:
              assignments = super(
                  AdamWithMultiLRSchedule, self
              ).apply_grad(grad, var)
      return assignments
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import re
  import gin
  from mesh_tensorflow import layers
  from mesh_tensorflow import ops_with_redefined_builtins
  import tensorflow.compat.v1
label_desc: Use image node_id 2 for calling the AdamWithMultiLRSchedule obj's underlying member method code with example usage: obj.apply_grad(grad, var) and returns: assignments, super