def adam(params, grads, lr, schedule, t_total,
         b1=0.9,
         b2=0.999,
         e=1e-8,
         weight_decay=1e-2,
         bias_l2=True,
         max_grad_norm=1.):
  """Custom Adam optimzizer for weight decay and learning rate schedule.

  Implementation adapted from https://github.com/openai/finetune-transformer-lm.

  Args:
    params: Parameters to be optimzed.
    grads: Gradients.
    lr: learning rate.
    schedule: Type of learning rate scheduling
    t_total: Total training steps.
    b1: beta_1.
    b2: beta_2.
    e: epsilon.
    weight_decay: Weight decay coefficient.
    bias_l2: Pose l2 penalty on bias parameters or not.
    max_grad_norm: Norm of gradient ot be clipped to.

  Returns:
    A list of update operations.
  """
  t = tf.train.get_global_step()
  tt = t + 1
  updates = [t.assign(tt)]
  if max_grad_norm > 0:
    grads, _ = tf.clip_by_global_norm(grads, max_grad_norm)
  for p, g in zip(params, grads):
    if p is None or g is None:
      print("can't train", p.name, g)
    else:
      if isinstance(g, tf.IndexedSlices):
        g = tf.convert_to_tensor(g)

      # past 1st moment vector; same shape as p.
      m = tf.Variable(p * 0., dtype=tf.float32, trainable=False)

      # past 2nd moment vector; same shape as p.
      v = tf.Variable(p * 0., dtype=tf.float32, trainable=False)
      lrt = lr * tf.sqrt(1 - b2**(tf.cast(tt, tf.float32))) / \
          (1 - b1**(tf.cast(tt, tf.float32)))
      lrt *= schedule(tf.cast(t, tf.float32)/t_total)

      # new 1st moment vector; same shape as p.
      mt = b1 * m + (1 - b1) * g

      # new 2nd moment vector; same shape as p.
      vt = b2 * v + (1 - b2) * g * g

      if (len(p.get_shape()) > 1 or bias_l2) and weight_decay > 0:
        pt = p - lrt * (mt / (tf.sqrt(vt) + e) + weight_decay * p)
      else:
        pt = p - lrt * (mt / (tf.sqrt(vt) + e))
      updates.extend([m.assign(mt), v.assign(vt), p.assign(pt)])
  return tf.group(*updates)
Custom Adam optimzizer for weight decay and learning rate schedule.

  Implementation adapted from https://github.com/openai/finetune-transformer-lm.

  Args:
    params: Parameters to be optimzed.
    grads: Gradients.
    lr: learning rate.
    schedule: Type of learning rate scheduling
    t_total: Total training steps.
    b1: beta_1.
    b2: beta_2.
    e: epsilon.
    weight_decay: Weight decay coefficient.
    bias_l2: Pose l2 penalty on bias parameters or not.
    max_grad_norm: Norm of gradient ot be clipped to.

  Returns:
    A list of update operations.
  
language/labs/exemplar_decoding/models/adam.py