//  [python]
def create_optimizer(loss, init_lr, num_train_steps, num_warmup_steps, accum_steps=1, use_tpu=False, use_multi_cpu=0, use_gpu=False):
  """Creates an optimizer training op."""
  global_step = tf.compat.v1.train.get_or_create_global_step()

  learning_rate = tf.constant(value=init_lr, shape=[], dtype=tf.float32)

  # Implements linear decay of the learning rate.
  learning_rate = tf.compat.v1.train.polynomial_decay(
      learning_rate,
      global_step,
      num_train_steps,
      end_learning_rate=0.0,
      power=1.0,
      cycle=False)

  # Implements linear warmup. I.e., if global_step < num_warmup_steps, the
  # learning rate will be `global_step/num_warmup_steps * init_lr`.
  if num_warmup_steps:
    global_steps_int = tf.cast(global_step, tf.int32)
    warmup_steps_int = tf.constant(num_warmup_steps, dtype=tf.int32)

    global_steps_float = tf.cast(global_steps_int, tf.float32)
    warmup_steps_float = tf.cast(warmup_steps_int, tf.float32)

    warmup_percent_done = global_steps_float / warmup_steps_float
    warmup_learning_rate = init_lr * warmup_percent_done

    is_warmup = tf.cast(global_steps_int < warmup_steps_int, tf.float32)
    learning_rate = (
        (1.0 - is_warmup) * learning_rate + is_warmup * warmup_learning_rate)


  # It is recommended that you use this optimizer for fine tuning, since this
  # is how the model was trained (note that the Adam m/v variables are NOT
  # loaded from init_checkpoint.)
  use_itex_optimizer = False
  if use_gpu:
    try:
      import intel_extension_for_tensorflow as itex
      optimizer = itex.ops.AdamWithWeightDecayOptimizer(
          learning_rate=learning_rate,
          weight_decay_rate=0.01,
          beta_1=0.9,
          beta_2=0.999,
          epsilon=1e-6,
          exclude_from_weight_decay=["LayerNorm", "layer_norm", "bias"])
      use_itex_optimizer = True
    except:
      print("Load plugin Intel(R) Extension for Tensorflow* failed.\n")

  if not use_itex_optimizer:
    optimizer = AdamWeightDecayOptimizer(
        learning_rate=learning_rate,
        weight_decay_rate=0.01,
        beta_1=0.9,
        beta_2=0.999,
        epsilon=1e-6,
        exclude_from_weight_decay=["LayerNorm", "layer_norm", "bias"])

  if use_multi_cpu and (accum_steps == 1):
    import horovod.tensorflow as hvd
    optimizer = hvd.DistributedOptimizer(optimizer, sparse_as_dense=True)

  if use_tpu:
    optimizer = tf.compat.v1.tpu.CrossShardOptimizer(optimizer)

  tvars = tf.compat.v1.trainable_variables()

  #################################################################################
  # support for grad accumulation
  #################################################################################
  if accum_steps > 1 :
    #tf.compat.v1.logging.info("Accumulation Steps....")
    grads_and_vars = optimizer.compute_gradients(loss * 1.0 / accum_steps, tvars)

    current_step = tf.compat.v1.get_variable(name="current_step", shape=[], dtype=tf.int32, 
                                   trainable=False,
                                   initializer=tf.zeros_initializer)
    accum_vars = [tf.compat.v1.get_variable(
          name=tvar.name.split(":")[0] + "/agrads",
          shape=tvar.shape.as_list(),
          dtype=tf.float32,
          trainable=False,
          initializer=tf.zeros_initializer()) for tvar in tvars]

    apply_grads = tf.cast(tf.math.equal(current_step % accum_steps, 0), dtype=tf.bool)
    current_step = tf.cond(apply_grads, 
                           lambda:current_step.assign(tf.ones_like(current_step)), 
                           lambda:current_step.assign_add(1))
                           #lambda:inc_current_step(current_step, "Apply Grads:"), 
                           #lambda:inc_current_step(current_step, "Step:"))

    grads_and_vars_and_accums = [(gv[0],gv[1],accum_vars[i]) for i, gv in enumerate(grads_and_vars) if gv[0] is not None]
    grads, tvars, accum_vars = list(zip(*grads_and_vars_and_accums))

    (cgrads, _) = tf.clip_by_global_norm(grads, clip_norm=1.0)

    #accum_vars = update_accum_vars(accum_vars, apply_grads, cgrads, current_step)
    accum_vars = tf.cond(apply_grads,
              lambda: [accum_vars[i].assign(grad) for i, grad in enumerate(cgrads)],
              lambda: [accum_vars[i].assign_add(grad) for i, grad in enumerate(cgrads)])

    def applyGrads(accum_vars, current_step):
          # if 1 or 0 MPI process, skip allreduce; otherwise do allreduce of the accum_var
          if use_multi_cpu > 1:
            import horovod.tensorflow as hvd
            accum_vars = [hvd.allreduce(tf.convert_to_tensor(accum_var)) if isinstance(accum_var, tf.IndexedSlices)
                                                              else hvd.allreduce(accum_var) for accum_var in accum_vars]
          #tf.compat.v1.logging.info("\t\t APPLYING GRADIENTS....:", global_step)
          return optimizer.apply_gradients(list(zip(accum_vars, tvars)), global_step=global_step)

    apply_step = tf.identity(tf.cast(tf.math.equal(current_step % accum_steps, 0), dtype=tf.bool), name="apply_step")
    update_op = tf.cond(apply_step, lambda: applyGrads(accum_vars, current_step), lambda: tf.no_op())

    new_global_step = tf.cond(apply_step, lambda: global_step+1, lambda: global_step)
    new_global_step = tf.identity(new_global_step, name='global_step_update')
    train_op = tf.group(update_op, [global_step.assign(new_global_step)])
  else :
    #grads = tf.gradients(ys=loss, xs=tvars)

    if use_multi_cpu:
      grads_and_vars = optimizer.compute_gradients(
          loss, tvars, gate_gradients=tf.compat.v1.train.Optimizer.GATE_NONE)
      grads = [grad for grad, var in grads_and_vars]
      tvars = [var for grad, var in grads_and_vars]
    else:
      grads = tf.gradients(loss, tvars)

    # This is how the model was pre-trained.
    (grads, _) = tf.clip_by_global_norm(grads, clip_norm=1.0)

    train_op = optimizer.apply_gradients(
        zip(grads, tvars), global_step=global_step)

    # Normally the global step update is done inside of `apply_gradients`.
    # However, `AdamWeightDecayOptimizer` doesn't do this. But if you use
    # a different optimizer, you should probably take this line out.
    if not use_itex_optimizer:
      new_global_step = global_step + 1
      new_global_step = tf.identity(new_global_step, name='global_step_update')
      train_op = tf.group(train_op, [global_step.assign(new_global_step)])

  return train_op