jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _is_sequence_right_padded(mask):
    """Check the mask tensor and see if it right padded.

    cuDNN uses the sequence length param to skip the tailing
    timestep. If the data is left padded, or not a strict right padding (has
    masked value in the middle of the sequence), then cuDNN won't work
    properly in those cases.

    Left padded data: [[False, False, True, True, True]].
    Right padded data: [[True, True, True, False, False]].
    Mixture of mask/unmasked data: [[True, False, True, False, False]].

    Note that for the mixed data example above, the actually data RNN should see
    are those 2 Trues (index 0 and 2), the index 1 False should be ignored and
    not pollute the internal states.

    Args:
        mask: the Boolean tensor with shape [batch, timestep]

    Returns:
        boolean scalar tensor, whether the mask is strictly right padded.
    """
    max_seq_length = tf.shape(mask)[1]
    count_of_true = tf.reduce_sum(tf.cast(mask, tf.int32), axis=1)
    right_padded_mask = tf.sequence_mask(count_of_true, maxlen=max_seq_length)
    return tf.reduce_all(
        tf.equal(
            tf.cast(mask, dtype="bool"),
            tf.cast(right_padded_mask, dtype="bool"),
        )
    )
//  [python]
def _compute_sequence_length_from_mask(mask, time_major):
    """Calculate the sequence length tensor (1-D) based on the masking tensor.

    The masking tensor is a 2D boolean tensor with shape [batch, timestep]. For
    any timestep that should be masked, the corresponding field will be False.
    Consider the following example:
      a = [[True, True, False, False],
           [True, True, True, False]]
    It is a (2, 4) tensor, and the corresponding sequence length result should
    be 1D tensor with value [2, 3]. Note that the masking tensor must be right
    padded that could be checked by, e.g., `is_sequence_right_padded()`.

    Args:
        mask: Boolean tensor with shape [batch, timestep] or [timestep, batch]
            if time_major=True.
        time_major: Boolean, which indicates whether the mask is time major or
            batch major.

    Returns:
        sequence_length: 1D int32 tensor.
    """
    timestep_index = 0 if time_major else 1
    return tf.reduce_sum(tf.to_int32(mask), axis=timestep_index)
//  [python]
def _compute_sequence_length_from_mask(mask, time_major):
    """Calculate the sequence length tensor (1-D) based on the masking tensor.

    The masking tensor is a 2D boolean tensor with shape [batch, timestep]. For
    any timestep that should be masked, the corresponding field will be False.
    Consider the following example:
      a = [[True, True, False, False],
           [True, True, True, False]]
    It is a (2, 4) tensor, and the corresponding sequence length result should
    be 1D tensor with value [2, 3]. Note that the masking tensor must be right
    padded that could be checked by, e.g., `is_sequence_right_padded()`.

    Args:
        mask: Boolean tensor with shape [batch, timestep] or [timestep, batch]
            if time_major=True.
        time_major: Boolean, which indicates whether the mask is time major or
            batch major.

    Returns:
        sequence_length: 1D int32 tensor.
    """
    timestep_index = 0 if time_major else 1
    return tf.reduce_sum(tf.cast(mask, tf.int32), axis=timestep_index)
//  [python]
def custom_loss(self, loss_obj, loss_scale=1., inputs=None, model=None):
    def loss_fn(y_true, y_pred):
      y_pred = t.cast(y_pred, tf.float32)
      loss = loss_obj(y_true, y_pred)
      loss = mt.reduce_over(loss)
      loss *= loss_scale
      return loss
    if inputs is not None or model is not None:
      return self.loss_wrapper(loss_fn, inputs, model)
    else:
      return loss_fn
//  [python]
def custom_loss(self, loss_obj, loss_scale=1., inputs=None, model=None):
    def loss_fn(y_true, y_pred):
      y_pred = tf.cast(y_pred, tf.float32)
      loss = loss_obj(y_true, y_pred)
      loss = mt.reduce_over(loss)
      loss *= loss_scale
      return loss
    if inputs is not None or model is not None:
      return self.loss_wrapper(loss_fn, inputs, model)
    else:
      return loss_fn
//  [python]
def monitor_emb(self, emb, name='emb', zero_ratio=False, training=True):
    if training:
      if self.step == 1 or not self.interval_steps or self.step % self.interval_steps == 0:
        if FLAGS.monitor_level > 1 and (not FLAGS.disable_monitor):
          res = {
            'l1': tf.reduce_mean(emb),
            'l2': tf.reduce_mean(emb * emb),
            'max': tf.reduce_max(emb),
            'min': tf.reduce_min(emb),
            'l2_0': tf.reduce_mean(emb[0] * emb[0]), 
            'l2_1': tf.reduce_mean(emb[1] * emb[1]),
            'l2_2': tf.reduce_mean(emb[2] * emb[2]), 
            'l2_-1': tf.reduce_mean(emb[-1] * emb[-1])
          }
          if zero_ratio:
            res['0ratio'] = tf.reduce_mean(tf.to_float(emb == 0))
            res['pos_ratio'] = tf.reduce_mean(tf.to_float(emb > 0))
          self.metrics_(gezi.dict_prefix(res, f'{name}/'))
//  [python]
def monitor_emb(self, emb, name='emb', zero_ratio=False, training=True):
    if training:
      if self.step == 1 or not self.interval_steps or self.step % self.interval_steps == 0:
        if FLAGS.monitor_level > 1 and (not FLAGS.disable_monitor):
          res = {
            'l1': tf.reduce_mean(emb),
            'l2': tf.reduce_mean(emb * emb),
            'max': tf.reduce_max(emb),
            'min': tf.reduce_min(emb),
            'l2_0': tf.reduce_mean(emb[0] * emb[0]), 
            'l2_1': tf.reduce_mean(emb[1] * emb[1]),
            'l2_2': tf.reduce_mean(emb[2] * emb[2]), 
            'l2_-1': tf.reduce_mean(emb[-1] * emb[-1])
          }
          if zero_ratio:
            res['0ratio'] = tf.reduce_mean(tf.cast(emb == 0, tf.float32))
            res['pos_ratio'] = tf.reduce_mean(tf.cast(emb > 0, tf.float32))
          self.metrics_(gezi.dict_prefix(res, f'{name}/'))
//  [python]
def fp16(x):
  dtype = tf.float16 if not gezi.get('tpu') else tf.bfloat16
  return tf.cast(x, dtype=dtype)
//  [python]
def float16_scope():
  """Scope class for float16."""

  def _custom_getter(getter, *args, **kwargs):
    """Returns a custom getter that methods must be called under."""
    cast_to_float16 = False
    requested_dtype = kwargs['dtype']
    if requested_dtype == tf.float16:
      kwargs['dtype'] = tf.float32
      cast_to_float16 = True
    var = getter(*args, **kwargs)
    if cast_to_float16:
      var = tf.cast(var, tf.float16)
    return var

  with tf.variable_scope('', custom_getter=_custom_getter) as varscope:
    yield varscope
//  [python]
def save_model(model, path, fp16=False):
  tf.keras.mixed_precision.experimental.set_policy('float32')
  gezi.try_mkdir(os.path.dirname(path))
  try:
    model = model.get_model()
  except Exception:
    pass
  try:
    info = gezi.get('info')
    optimizer = info['optimizer']
    optimizer._set_hyper('learning_rate', 0.001)
    model.compile(optimizer=optimizer,
                  loss=info['loss_fn'],
                  metrics=info['metrics'])
  except Exception:
    pass
  if os.path.isdir(path):
    path = os.path.join(path, 'model.h5')

  try:
    if not fp16:
      model.save(path)
    else:
      import pickle
      net = model.to_json()
      weights = [tf.cast(x, tf.float16) for x in model.weights]
      m = {'net': net, 'weights': weights}
      pickle.dump(m, open(path, 'wb'))
    model_size = gezi.get_size(path)
    logging.info(f'melt.save_model to {path} with model_size: {model_size}M')
    run = gezi.get('wandb_run')
    if run:
      wandb.log({'Others/model_size': model_size})
  except Exception as e:
    logging.warning(f'melt.save_model to {path} fail')
    logging.warning(e)
//  [python]
def rm_nan_inf(x, val=0.0):
  """Remove nan and infinity.

  Args:
    x: any `Tensor` of any type.
    val: value to replace nan and infinity with.

  Returns:
    a `Tensor` with nan and infinity removed.
  """
  cond = tf.math.logical_or(tf.math.is_nan(x), tf.math.is_inf(x))
  val = tf.to_float(val)
  x = tf.where(cond, val, x)
  return x
//  [python]
def rm_nan_inf(x, val=0.0):
  """Remove nan and infinity.

  Args:
    x: any `Tensor` of any type.
    val: value to replace nan and infinity with.

  Returns:
    a `Tensor` with nan and infinity removed.
  """
  cond = tf.math.logical_or(tf.math.is_nan(x), tf.math.is_inf(x))
  val = tf.cast(val, dtype=x.dtype)
  x = tf.where(cond, val, x)
  return x
//  [python]
def rm_nan(x, val=0.0):
  """Remove nan and infinity.

  Args:
    x: any `Tensor` of any type.
    val: value to replace nan.

  Returns:
    a `Tensor` with nan removed.
  """
  cond = tf.math.is_nan(x)
  val = t.cast(val, dtype=x.dtype)
  x = tf.where(cond, val, x)
  return x
//  [python]
def rm_nan(x, val=0.0):
  """Remove nan and infinity.

  Args:
    x: any `Tensor` of any type.
    val: value to replace nan.

  Returns:
    a `Tensor` with nan removed.
  """
  cond = tf.math.is_nan(x)
  val = tf.cast(val, dtype=x.dtype)
  x = tf.where(cond, val, x)
  return x
//  [python]
def rpn_losses_giou(
    gt_objectness_logits,
    gt_anchor_deltas,
    pred_objectness_logits,
    pred_anchor_deltas,
):
    reg_loss_sum = 1.0 - odl.giou(pred_anchor_deltas, gt_anchor_deltas)
    localization_loss = tf.reduce_sum(reg_loss_sum)

    objectness_loss = tf.losses.sigmoid_cross_entropy(
        logits=tf.expand_dims(pred_objectness_logits,1),
        multi_class_labels=tf.cast(tf.expand_dims(gt_objectness_logits,axis=1),tf.float32),
        reduction=tf.losses.Reduction.SUM,
        loss_collection=None
    )
    return objectness_loss, localization_loss
//  [python]
def build_model_with_precision(pp, mm, ii, *args, **kwargs):
  """Build model with its inputs/params for a specified precision context.
  This is highly specific to this codebase, and not intended to be general API.
  Advanced users only. DO NOT use it if you don't know what it does.
  NOTE: short argument names are intended to avoid conficts with kwargs.
  Args:
    pp: A string, precision policy name, such as "mixed_float16".
    mm: A function, for rmodel builder.
    ii: A tensor, for model inputs.
    tt: A bool, If true, it is for training; otherwise, it is for eval.
    *args: A list of model arguments.
    **kwargs: A dict, extra model parameters.
  Returns:
    the output of mm model.
  """
  tt = True
  if pp == 'mixed_bfloat16':
    set_precision_policy(pp)
    inputs = t.cast(ii, tf.bfloat16)
    with tf.tpu.bfloat16_scope():
      outputs = mm(inputs, *args, **kwargs)
    set_precision_policy('float32')
  elif pp == 'mixed_float16':
    set_precision_policy(pp, loss_scale=tt)
    inputs = t.cast(ii, tf.float16)
    with float16_scope():
      outputs = mm(inputs, *args, **kwargs)
    set_precision_policy('float32')
  elif not pp or pp == 'float32':
    outputs = mm(ii, *args, **kwargs)
  else:
    raise ValueError('Unknow precision name {}'.format(pp))

  # Users are responsible to convert the dtype of all outputs.
  return outputs
//  [python]
def build_model_with_precision(pp, mm, ii, *args, **kwargs):
  """Build model with its inputs/params for a specified precision context.
  This is highly specific to this codebase, and not intended to be general API.
  Advanced users only. DO NOT use it if you don't know what it does.
  NOTE: short argument names are intended to avoid conficts with kwargs.
  Args:
    pp: A string, precision policy name, such as "mixed_float16".
    mm: A function, for rmodel builder.
    ii: A tensor, for model inputs.
    tt: A bool, If true, it is for training; otherwise, it is for eval.
    *args: A list of model arguments.
    **kwargs: A dict, extra model parameters.
  Returns:
    the output of mm model.
  """
  tt = True
  if pp == 'mixed_bfloat16':
    set_precision_policy(pp)
    inputs = tf.cast(ii, tf.bfloat16)
    with tf.tpu.bfloat16_scope():
      outputs = mm(inputs, *args, **kwargs)
    set_precision_policy('float32')
  elif pp == 'mixed_float16':
    set_precision_policy(pp, loss_scale=tt)
    inputs = tf.cast(ii, tf.float16)
    with float16_scope():
      outputs = mm(inputs, *args, **kwargs)
    set_precision_policy('float32')
  elif not pp or pp == 'float32':
    outputs = mm(ii, *args, **kwargs)
  else:
    raise ValueError('Unknow precision name {}'.format(pp))

  # Users are responsible to convert the dtype of all outputs.
  return outputs
//  [python]
def get_nll(self, seq_logits, doc_scores, target, reduce_loss=False, epsilon=0.0, n_docs=None):
        n_docs = n_docs if n_docs is not None else self.config.n_docs
        # shift tokens left (from original Pytorch's version)

        target = tf.concat(
            [target[:, 1:], tf.fill([target.shape[0], 1], t.cast(self.config.generator.pad_token_id, target.dtype))],
            axis=1,
        )
        rag_logprobs = self.marginalize(seq_logits, doc_scores, n_docs)
        loss = self.hf_compute_loss(target, rag_logprobs, from_logits=True, reduce_loss=reduce_loss)

        return loss
//  [python]
def get_nll(self, seq_logits, doc_scores, target, reduce_loss=False, epsilon=0.0, n_docs=None):
        n_docs = n_docs if n_docs is not None else self.config.n_docs
        # shift tokens left (from original Pytorch's version)

        target = tf.concat(
            [target[:, 1:], tf.fill([target.shape[0], 1], tf.cast(self.config.generator.pad_token_id, target.dtype))],
            axis=1,
        )
        rag_logprobs = self.marginalize(seq_logits, doc_scores, n_docs)
        loss = self.hf_compute_loss(target, rag_logprobs, from_logits=True, reduce_loss=reduce_loss)

        return loss
//  [python]
def _cat_and_pad(tensors, pad_token_id):
        # used by generate(): tensors is a (batched) list of (candidates, len); len is varied across batch

        # Initialize padded tensor with shape ( all_candidates , max_candidate_length ),
        # where all_candidates counted from all inputs
        new_shape = sum([t.shape[0] for t in tensors]), max([t.shape[1] for t in tensors])
        output = tf.fill(new_shape, pad_token_id)

        # Normal tensor doesn't support slice assignment, so we need tf.Variable
        output = tf.Variable(output)

        # Assign, and then convert back to tensor
        ind = 0
        for t in tensors:
            output[ind : ind + t.shape[0], : t.shape[1]].assign(t)
            ind += t.shape[0]

        output = tf.convert_to_tensor(output)
        return t.cast(output, tensors[0][0][0].dtype)
//  [python]
def _cat_and_pad(tensors, pad_token_id):
        # used by generate(): tensors is a (batched) list of (candidates, len); len is varied across batch

        # Initialize padded tensor with shape ( all_candidates , max_candidate_length ),
        # where all_candidates counted from all inputs
        new_shape = sum([t.shape[0] for t in tensors]), max([t.shape[1] for t in tensors])
        output = tf.fill(new_shape, pad_token_id)

        # Normal tensor doesn't support slice assignment, so we need tf.Variable
        output = tf.Variable(output)

        # Assign, and then convert back to tensor
        ind = 0
        for t in tensors:
            output[ind : ind + t.shape[0], : t.shape[1]].assign(t)
            ind += t.shape[0]

        output = tf.convert_to_tensor(output)
        return tf.cast(output, tensors[0][0][0].dtype)
//  [python]
def shift_tokens_right(self, input_ids, start_token_id=None):
        """Shift input ids one token to the right, and pad with start_token_id"""

        if start_token_id is None:
            start_token_id = self.generator.config.decoder_start_token_id
            assert start_token_id is not None, (
                "self.generator.config.decoder_start_token_id has to be defined. In Rag we commonly use Bart as"
                " generator, see Bart docs for more information"
            )

        pad_token_id = self.generator.config.pad_token_id
        assert pad_token_id is not None, "self.model.config.pad_token_id has to be defined."

        start_tokens = tf.fill((shape_list(input_ids)[0], 1), tf.to_int32(start_token_id))
        shifted_input_ids = tf.concat([start_tokens, input_ids[:, :-1]], -1)

        # replace possible -100 values in labels by `pad_token_id`
        shifted_input_ids = tf.where(
            shifted_input_ids == -100,
            tf.fill(shape_list(shifted_input_ids), tf.to_int32(pad_token_id)),
            shifted_input_ids,
        )

        # "Verify that `labels` has only positive values and -100"
        assert_gte0 = tf.debugging.assert_greater_equal(shifted_input_ids, tf.to_int32(0))

        # Make sure the assertion op is called by wrapping the result in an identity no-op
        with tf.control_dependencies([assert_gte0]):
            shifted_input_ids = tf.identity(shifted_input_ids)

        return shifted_input_ids
//  [python]
def shift_tokens_right(self, input_ids, start_token_id=None):
        """Shift input ids one token to the right, and pad with start_token_id"""

        if start_token_id is None:
            start_token_id = self.generator.config.decoder_start_token_id
            assert start_token_id is not None, (
                "self.generator.config.decoder_start_token_id has to be defined. In Rag we commonly use Bart as"
                " generator, see Bart docs for more information"
            )

        pad_token_id = self.generator.config.pad_token_id
        assert pad_token_id is not None, "self.model.config.pad_token_id has to be defined."

        start_tokens = tf.fill((shape_list(input_ids)[0], 1), tf.cast(start_token_id, input_ids.dtype))
        shifted_input_ids = tf.concat([start_tokens, input_ids[:, :-1]], -1)

        # replace possible -100 values in labels by `pad_token_id`
        shifted_input_ids = tf.where(
            shifted_input_ids == -100,
            tf.fill(shape_list(shifted_input_ids), tf.cast(pad_token_id, input_ids.dtype)),
            shifted_input_ids,
        )

        # "Verify that `labels` has only positive values and -100"
        assert_gte0 = tf.debugging.assert_greater_equal(shifted_input_ids, tf.cast(0, shifted_input_ids.dtype))

        # Make sure the assertion op is called by wrapping the result in an identity no-op
        with tf.control_dependencies([assert_gte0]):
            shifted_input_ids = tf.identity(shifted_input_ids)

        return shifted_input_ids
//  [python]
def _get_image_info(self, image):
        desired_size = t.convert_to_tensor(self._input_shape, dtype=t.float32)
        image_size = t.cast(t.shape(input=image)[0:2], t.float32)
        scaled_size = desired_size

        scale = t.minimum(scaled_size[0] / image_size[0], scaled_size[1] / image_size[1])
        scaled_size = t.round(image_size * scale)

        # Computes 2D image_scale.
        image_scale = scaled_size / image_size

        offset = t.zeros((2,), t.int32)

        image_info = t.stack([image_size, t.cast(desired_size, t.float32), image_scale, t.cast(offset, t.float32)])
        return image_info
//  [python]
def _get_image_info(self, image):
        desired_size = tf.convert_to_tensor(self._input_shape, dtype=tf.float32)
        image_size = tf.cast(tf.shape(input=image)[0:2], tf.float32)
        scaled_size = desired_size

        scale = tf.minimum(scaled_size[0] / image_size[0], scaled_size[1] / image_size[1])
        scaled_size = tf.round(image_size * scale)

        # Computes 2D image_scale.
        image_scale = scaled_size / image_size

        offset = tf.zeros((2,), tf.int32)

        image_info = tf.stack([image_size, tf.cast(desired_size, tf.float32), image_scale, tf.cast(offset, tf.float32)])
        return image_info
//  [python]
def _parse_predict_data(self, data):
        """Parses data for prediction"""
        image_data = data["image"]
        image_shape = tf.shape(input=image_data)[0:2]

        # needed only for eval
        image_info = self._get_image_info(image_data)

        # image preprocessing
        image_data = tf.py_function(self._preprocess_predict_image, [image_data], Tout=tf.float32)
        image_data.set_shape([None, None, 3])

        labels = {
            "image_info": image_info,
        }

        # Converts boxes from normalized coordinates to pixel coordinates.
        boxes = box_utils.denormalize_boxes(data["groundtruth_boxes"], image_shape)
        groundtruths = {
            "source_id": data["source_id"],
            "num_detections": tf.squeeze(tf.shape(data["groundtruth_classes"])),
            "boxes": boxes,
            "classes": data["groundtruth_classes"],
            "areas": data["groundtruth_area"],
            "is_crowds": tf.to_int32(data["groundtruth_is_crowd"]),
        }
        groundtruths["source_id"] = dataloader_utils.process_source_id(groundtruths["source_id"])
        groundtruths = dataloader_utils.pad_groundtruths_to_fixed_size(groundtruths, self._max_num_instances)
        labels.update(groundtruths)

        return image_data, labels
//  [python]
def _parse_predict_data(self, data):
        """Parses data for prediction"""
        image_data = data["image"]
        image_shape = tf.shape(input=image_data)[0:2]

        # needed only for eval
        image_info = self._get_image_info(image_data)

        # image preprocessing
        image_data = tf.py_function(self._preprocess_predict_image, [image_data], Tout=tf.float32)
        image_data.set_shape([None, None, 3])

        labels = {
            "image_info": image_info,
        }

        # Converts boxes from normalized coordinates to pixel coordinates.
        boxes = box_utils.denormalize_boxes(data["groundtruth_boxes"], image_shape)
        groundtruths = {
            "source_id": data["source_id"],
            "num_detections": tf.squeeze(tf.shape(data["groundtruth_classes"])),
            "boxes": boxes,
            "classes": data["groundtruth_classes"],
            "areas": data["groundtruth_area"],
            "is_crowds": tf.cast(data["groundtruth_is_crowd"], tf.int32),
        }
        groundtruths["source_id"] = dataloader_utils.process_source_id(groundtruths["source_id"])
        groundtruths = dataloader_utils.pad_groundtruths_to_fixed_size(groundtruths, self._max_num_instances)
        labels.update(groundtruths)

        return image_data, labels
//  [python]
def pixel_center_grid(height, width):
  """Produce a grid of (x,y) texture-coordinate pairs of pixel centers.

  Args:
    height: (integer) height, not a tensor
    width: (integer) width, not a tensor

  Returns:
    A tensor of shape [height, width, 2] where each entry gives the (x,y)
    texture coordinates of the corresponding pixel center. For example, for
    pixel_center_grid(2, 3) the result is:
       [[[1/6, 1/4], [3/6, 1/4], [5/6, 1/4]],
        [[1/6, 3/4], [3/6, 3/4], [5/6, 3/4]]]
  """
  height_float = t.cast(height, dtype=t.float32)
  width_float = t.cast(width, dtype=t.float32)
  ys = tf.linspace(0.5 / height_float, 1.0 - 0.5 / height_float, height)
  xs = tf.linspace(0.5 / width_float, 1.0 - 0.5 / width_float, width)
  xs, ys = tf.meshgrid(xs, ys)
  grid = tf.stack([xs, ys], axis=-1)
  assert grid.shape.as_list() == [height, width, 2]
  return grid
//  [python]
def pixel_center_grid(height, width):
  """Produce a grid of (x,y) texture-coordinate pairs of pixel centers.

  Args:
    height: (integer) height, not a tensor
    width: (integer) width, not a tensor

  Returns:
    A tensor of shape [height, width, 2] where each entry gives the (x,y)
    texture coordinates of the corresponding pixel center. For example, for
    pixel_center_grid(2, 3) the result is:
       [[[1/6, 1/4], [3/6, 1/4], [5/6, 1/4]],
        [[1/6, 3/4], [3/6, 3/4], [5/6, 3/4]]]
  """
  height_float = tf.cast(height, dtype=tf.float32)
  width_float = tf.cast(width, dtype=tf.float32)
  ys = tf.linspace(0.5 / height_float, 1.0 - 0.5 / height_float, height)
  xs = tf.linspace(0.5 / width_float, 1.0 - 0.5 / width_float, width)
  xs, ys = tf.meshgrid(xs, ys)
  grid = tf.stack([xs, ys], axis=-1)
  assert grid.shape.as_list() == [height, width, 2]
  return grid
//  [python]
def __build_masks(self):
    """Build pruning masks for all the maskable variables.

    Returns:
    * masks: list of pruning masks
    * pr_all: placeholder of all the pruning ratios
    * pr_assign_op: pruning ratios' assign operation
    """

    # create masks & pruning ratios
    masks, prune_ratios = [], []
    for var_full, var_prnd in zip(self.vars_full['maskable'], self.vars_prnd['maskable']):
      name = var_prnd.name.replace(':0', '_prune_ratio')
      prune_ratio = tf.get_variable(name, shape=[], trainable=False)
      mask_thres = tf.contrib.distributions.percentile(tf.abs(var_full), prune_ratio * 100)
      mask = tf.to_float(tf.abs(var_full) > mask_thres)
      masks += [mask]
      prune_ratios += [prune_ratio]

    # create a placeholder to pass pruning ratios to the TF tensor
    pr_all = tf.placeholder(tf.float32, shape=(len(prune_ratios)), name='pr_all')
    pr_all_split = tf.split(pr_all, len(prune_ratios))
    pr_assign_op = tf.group([prune_ratio.assign(tf.reshape(pr, []))
                             for prune_ratio, pr in zip(prune_ratios, pr_all_split)])

    return masks, pr_all, pr_assign_op
//  [python]
def __build_masks(self):
    """Build pruning masks for all the maskable variables.

    Returns:
    * masks: list of pruning masks
    * pr_all: placeholder of all the pruning ratios
    * pr_assign_op: pruning ratios' assign operation
    """

    # create masks & pruning ratios
    masks, prune_ratios = [], []
    for var_full, var_prnd in zip(self.vars_full['maskable'], self.vars_prnd['maskable']):
      name = var_prnd.name.replace(':0', '_prune_ratio')
      prune_ratio = tf.get_variable(name, shape=[], trainable=False)
      mask_thres = tf.contrib.distributions.percentile(tf.abs(var_full), prune_ratio * 100)
      mask = tf.cast(tf.abs(var_full) > mask_thres, tf.float32)
      masks += [mask]
      prune_ratios += [prune_ratio]

    # create a placeholder to pass pruning ratios to the TF tensor
    pr_all = tf.placeholder(tf.float32, shape=(len(prune_ratios)), name='pr_all')
    pr_all_split = tf.split(pr_all, len(prune_ratios))
    pr_assign_op = tf.group([prune_ratio.assign(tf.reshape(pr, []))
                             for prune_ratio, pr in zip(prune_ratios, pr_all_split)])

    return masks, pr_all, pr_assign_op
//  [python]
def _prune_weight(self, op, weight_tensor, scope):
        with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
            weight_abs_tensor = tf.abs(weight_tensor)
            weight_sorted_abs_tensor = tf.contrib.framework.sort(tf.reshape(weight_abs_tensor, [-1]))
            num = weight_sorted_abs_tensor.shape.num_elements() * self._config_dict[op.name] / self._total_pruning_iterations * tf.cast(self._step, tf.float32)
            num = tf.to_int64(num)
            threshold = weight_sorted_abs_tensor[num]
            mask = tf.to_float(weight_abs_tensor > threshold)
            self._mask_dict[op.name] = mask

            if self._debug_info:
                pruning_ratio = 1.0 - tf.reduce_mean(mask)
                print_op = tf.print("pruning iter:", self._step, "/", self._total_pruning_iterations, op.name, "pruning_ratio:", pruning_ratio)

                with tf.control_dependencies([print_op]):
                    pruned_weight_tensor = weight_tensor * mask
            else:
                pruned_weight_tensor = weight_tensor * mask

            weight_variable_name = self._get_weight_variable_name(weight_tensor)
            weight_variable = self._get_variable_by_name(weight_variable_name)
            new_weight_tensor_assign = tf.assign(weight_variable, pruned_weight_tensor)

            tf.add_to_collection("update_pruning_masks", new_weight_tensor_assign)

            return new_weight_tensor_assign
//  [python]
def _prune_weight(self, op, weight_tensor, scope):
        with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
            weight_abs_tensor = tf.abs(weight_tensor)
            weight_sorted_abs_tensor = tf.contrib.framework.sort(tf.reshape(weight_abs_tensor, [-1]))
            num = weight_sorted_abs_tensor.shape.num_elements() * self._config_dict[op.name] / self._total_pruning_iterations * tf.cast(self._step, tf.float32)
            num = tf.cast(num, tf.int64)
            threshold = weight_sorted_abs_tensor[num]
            mask = tf.cast(weight_abs_tensor > threshold, tf.float32)
            self._mask_dict[op.name] = mask

            if self._debug_info:
                pruning_ratio = 1.0 - tf.reduce_mean(mask)
                print_op = tf.print("pruning iter:", self._step, "/", self._total_pruning_iterations, op.name, "pruning_ratio:", pruning_ratio)

                with tf.control_dependencies([print_op]):
                    pruned_weight_tensor = weight_tensor * mask
            else:
                pruned_weight_tensor = weight_tensor * mask

            weight_variable_name = self._get_weight_variable_name(weight_tensor)
            weight_variable = self._get_variable_by_name(weight_variable_name)
            new_weight_tensor_assign = tf.assign(weight_variable, pruned_weight_tensor)

            tf.add_to_collection("update_pruning_masks", new_weight_tensor_assign)

            return new_weight_tensor_assign
//  [python]
def sample_image(image, coords, clamp=True):
  """Sample points from an image, using bilinear filtering.

  Args:
    image: [B0, ..., Bn-1, height, width, channels] image data
    coords: [B0, ..., Bn-1, ..., 2] (x,y) texture coordinates
    clamp: if True, coordinates are clamped to the coordinates of the corner
      pixels -- i.e. minimum value 0.5/width, 0.5/height and maximum value
      1.0-0.5/width or 1.0-0.5/height. This is equivalent to extending the image
      in all directions by copying its edge pixels. If False, sampling values
      outside the image will return 0 values.

  Returns:
    [B0, ..., Bn-1, ..., channels] image data, in which each value is sampled
    with bilinear interpolation from the image at position indicated by the
    (x,y) texture coordinates. The image and coords parameters must have
    matching batch dimensions B0, ..., Bn-1.

  Raises:
    ValueError: if shapes are incompatible.
  """
  check_input_shape('coords', coords, -1, 2)
  tfshape = tf.shape(image)[-3:-1]
  height = tf.cast(tfshape[0], dtype=tf.float32)
  width = tf.cast(tfshape[1], dtype=tf.float32)
  if clamp:
    coords = clip_texture_coords_to_corner_pixels(coords, height, width)

  # Resampler expects coordinates where (0,0) is the center of the top-left
  # pixel and (width-1, height-1) is the center of the bottom-right pixel.
  pixel_coords = coords * [width, height] - 0.5

  # tfa_image.resampler only works with exactly one batch dimension, i.e. it
  # expects image to be [batch, height, width, channels] and pixel_coords to be
  # [batch, ..., 2]. So we need to reshape, perform the resampling, and then
  # reshape back to what we had.
  batch_dims = len(image.shape.as_list()) - 3
  assert (image.shape.as_list()[:batch_dims] == pixel_coords.shape.as_list()
          [:batch_dims])

  batched_image, _ = utils.flatten_batch(image, batch_dims)
  batched_coords, unflatten_coords = utils.flatten_batch(
      pixel_coords, batch_dims)
  resampled = tfa_image.resampler(batched_image, batched_coords)

  # Convert back to the right shape to return
  resampled = unflatten_coords(resampled)
  return resampled
//  [python]
def _localization_loss(pred_locs, gt_locs, gt_labels, num_matched_boxes):
  """Computes the localization loss.

  Computes the localization loss using smooth l1 loss.
  Args:
    pred_locs: a dict from index to tensor of predicted locations. The shape
      of each tensor is [batch_size, num_anchors, 4].
    gt_locs: a list of tensors representing box regression targets in
      [batch_size, num_anchors, 4].
    gt_labels: a list of tensors that represents the classification groundtruth
      targets. The shape is [batch_size, num_anchors, 1].
    num_matched_boxes: the number of anchors that are matched to a groundtruth
      targets, used as the loss normalizater. The shape is [batch_size].
  Returns:
    box_loss: a float32 representing total box regression loss.
  """
  keys = sorted(pred_locs.keys())
  box_loss = 0
  for i, k in enumerate(keys):
    gt_label = gt_labels[i]
    gt_loc = gt_locs[i]
    pred_loc = tf.reshape(pred_locs[k], gt_loc.shape)
    mask = tf.greater(gt_label, 0)
    float_mask = tf.to_float(mask)

    smooth_l1 = tf.reduce_sum(
        tf.losses.huber_loss(
            gt_loc, pred_loc, reduction=tf.losses.Reduction.NONE),
        axis=-1)
    smooth_l1 = tf.multiply(smooth_l1, float_mask)
    box_loss = box_loss + tf.reduce_sum(
        smooth_l1, axis=list(range(1, smooth_l1.shape.ndims)))

  # TODO(taylorrobie): Confirm that normalizing by the number of boxes matches
  # reference
  return tf.reduce_mean(box_loss / num_matched_boxes)
//  [python]
def _localization_loss(pred_locs, gt_locs, gt_labels, num_matched_boxes):
  """Computes the localization loss.

  Computes the localization loss using smooth l1 loss.
  Args:
    pred_locs: a dict from index to tensor of predicted locations. The shape
      of each tensor is [batch_size, num_anchors, 4].
    gt_locs: a list of tensors representing box regression targets in
      [batch_size, num_anchors, 4].
    gt_labels: a list of tensors that represents the classification groundtruth
      targets. The shape is [batch_size, num_anchors, 1].
    num_matched_boxes: the number of anchors that are matched to a groundtruth
      targets, used as the loss normalizater. The shape is [batch_size].
  Returns:
    box_loss: a float32 representing total box regression loss.
  """
  keys = sorted(pred_locs.keys())
  box_loss = 0
  for i, k in enumerate(keys):
    gt_label = gt_labels[i]
    gt_loc = gt_locs[i]
    pred_loc = tf.reshape(pred_locs[k], gt_loc.shape)
    mask = tf.greater(gt_label, 0)
    float_mask = tf.cast(mask, tf.float32)

    smooth_l1 = tf.reduce_sum(
        tf.losses.huber_loss(
            gt_loc, pred_loc, reduction=tf.losses.Reduction.NONE),
        axis=-1)
    smooth_l1 = tf.multiply(smooth_l1, float_mask)
    box_loss = box_loss + tf.reduce_sum(
        smooth_l1, axis=list(range(1, smooth_l1.shape.ndims)))

  # TODO(taylorrobie): Confirm that normalizing by the number of boxes matches
  # reference
  return tf.reduce_mean(box_loss / num_matched_boxes)
//  [python]
def learning_rate_schedule(params, global_step):
  """Handles learning rate scaling, linear warmup, and learning rate decay.

  Args:
    params: A dictionary that defines hyperparameters of model.
    global_step: A tensor representing current global step.

  Returns:
    A tensor representing current learning rate.
  """
  base_learning_rate = params['base_learning_rate']
  lr_warmup_step = params['lr_warmup_step']
  cos_decay_step = params['cos_decay_step']
  batch_size = params['batch_size']
  scaling_factor = get_rank_size() * batch_size / ssd_constants.DEFAULT_BATCH_SIZE
  adjusted_learning_rate = base_learning_rate * scaling_factor
  learning_rate = (tf.to_float(global_step) /
                   lr_warmup_step) * adjusted_learning_rate

  learning_rate = tf.where(global_step < lr_warmup_step, learning_rate,
                           tf.train.cosine_decay(adjusted_learning_rate, global_step, cos_decay_step, alpha=0.01))

  return learning_rate
//  [python]
def learning_rate_schedule(params, global_step):
  """Handles learning rate scaling, linear warmup, and learning rate decay.

  Args:
    params: A dictionary that defines hyperparameters of model.
    global_step: A tensor representing current global step.

  Returns:
    A tensor representing current learning rate.
  """
  base_learning_rate = params['base_learning_rate']
  lr_warmup_step = params['lr_warmup_step']
  cos_decay_step = params['cos_decay_step']
  batch_size = params['batch_size']
  scaling_factor = get_rank_size() * batch_size / ssd_constants.DEFAULT_BATCH_SIZE
  adjusted_learning_rate = base_learning_rate * scaling_factor
  learning_rate = (tf.cast(global_step, dtype=tf.float32) /
                   lr_warmup_step) * adjusted_learning_rate

  learning_rate = tf.where(global_step < lr_warmup_step, learning_rate,
                           tf.train.cosine_decay(adjusted_learning_rate, global_step, cos_decay_step, alpha=0.01))

  return learning_rate
//  [python]
def _lower_tail(self, tail_mass):
    logits = tf.math.log(
        tf.cast(tail_mass / 2 / (1. - tail_mass / 2), self.dtype))
    return helpers.estimate_tails(
        self._logits_cumulative, logits, self.batch_shape_tensor(), self.dtype)
//  [python]
def _upper_tail(self, tail_mass):
    logits = -tf.math.log(
        tf.cast(tail_mass / 2 / (1. - tail_mass / 2), self.dtype))
    return helpers.estimate_tails(
        self._logits_cumulative, logits, self.batch_shape_tensor(), self.dtype)
//  [python]
def random_substr(str_tensor, max_words):
  """Select random substring if the input has more than max_words."""
  word_batch_r = tf.strings.split(str_tensor)
  row_splits = word_batch_r.row_splits
  words = word_batch_r.values
  start_idx = row_splits[:-1]
  end_idx = row_splits[1:]
  words_per_example = end_idx - start_idx
  ones = tf.ones_like(end_idx)
  max_val = tf.maximum(ones, words_per_example - max_words)
  max_words_batch = tf.reduce_max(words_per_example)
  rnd = tf.random.uniform(
      tf.shape(start_idx), minval=0, maxval=max_words_batch, dtype=tf.int64)
  off_start_idx = tf.math.floormod(rnd, max_val)
  new_words_per_example = tf.where(
      tf.equal(max_val, 1), words_per_example, ones * max_words)
  new_start_idx = start_idx + off_start_idx
  new_end_idx = new_start_idx + new_words_per_example
  indices = tf.expand_dims(tf.range(tf.size(words), dtype=tf.int64), axis=0)
  within_limit = tf.logical_and(
      tf.greater_equal(indices, tf.expand_dims(new_start_idx, axis=1)),
      tf.less(indices, tf.expand_dims(new_end_idx, axis=1)))
  keep_indices = tf.reduce_any(within_limit, axis=0)
  keep_indices = tf.to_int32(keep_indices)
  _, selected_words = tf.dynamic_partition(words, keep_indices, 2)
  row_splits = tf.math.cumsum(new_words_per_example)
  row_splits = tf.concat([[0], row_splits], axis=0)
  new_tensor = tf.RaggedTensor.from_row_splits(
      values=selected_words, row_splits=row_splits)
  return tf.strings.reduce_join(new_tensor, axis=1, separator=" ")
//  [python]
def random_substr(str_tensor, max_words):
  """Select random substring if the input has more than max_words."""
  word_batch_r = tf.strings.split(str_tensor)
  row_splits = word_batch_r.row_splits
  words = word_batch_r.values
  start_idx = row_splits[:-1]
  end_idx = row_splits[1:]
  words_per_example = end_idx - start_idx
  ones = tf.ones_like(end_idx)
  max_val = tf.maximum(ones, words_per_example - max_words)
  max_words_batch = tf.reduce_max(words_per_example)
  rnd = tf.random.uniform(
      tf.shape(start_idx), minval=0, maxval=max_words_batch, dtype=tf.int64)
  off_start_idx = tf.math.floormod(rnd, max_val)
  new_words_per_example = tf.where(
      tf.equal(max_val, 1), words_per_example, ones * max_words)
  new_start_idx = start_idx + off_start_idx
  new_end_idx = new_start_idx + new_words_per_example
  indices = tf.expand_dims(tf.range(tf.size(words), dtype=tf.int64), axis=0)
  within_limit = tf.logical_and(
      tf.greater_equal(indices, tf.expand_dims(new_start_idx, axis=1)),
      tf.less(indices, tf.expand_dims(new_end_idx, axis=1)))
  keep_indices = tf.reduce_any(within_limit, axis=0)
  keep_indices = tf.cast(keep_indices, dtype=tf.int32)
  _, selected_words = tf.dynamic_partition(words, keep_indices, 2)
  row_splits = tf.math.cumsum(new_words_per_example)
  row_splits = tf.concat([[0], row_splits], axis=0)
  new_tensor = tf.RaggedTensor.from_row_splits(
      values=selected_words, row_splits=row_splits)
  return tf.strings.reduce_join(new_tensor, axis=1, separator=" ")
//  [python]
def _classification_loss(pred_labels, gt_labels, num_matched_boxes):
  """Computes the classification loss.

  Computes the classification loss with hard negative mining.
  Args:
    pred_labels: a dict from index to tensor of predicted class. The shape
      of the tensor is [batch_size, num_anchors, num_classes].
    gt_labels: a list of tensor that represents the classification groundtruth
      targets. The shape is [batch_size, num_anchors, 1].
    num_matched_boxes: the number of anchors that are matched to a groundtruth
      targets. This is used as the loss normalizater.
  Returns:
    box_loss: a float32 representing total box regression loss.
  """
  keys = sorted(pred_labels.keys())
  batch_size = gt_labels[0].shape[0]
  cross_entropy = []
  for i, k in enumerate(keys):
    gt_label = gt_labels[i]
    pred_label = tf.reshape(
        pred_labels[k],
        gt_label.get_shape().as_list() + [ssd_constants.NUM_CLASSES])
    cross_entropy.append(
        tf.reshape(
            _softmax_cross_entropy(pred_label, gt_label), [batch_size, -1]))


  # Put the rest of the loss computation on one device to avoid excessive
  # communication inside topk_mask with spatial partition
  #with tf.device(tf.contrib.tpu.core(0)):
  cross_entropy = tf.concat(cross_entropy, 1)
  gt_label = tf.concat([tf.reshape(l, [batch_size, -1]) for l in gt_labels],
                         1)
  mask = tf.greater(gt_label, 0)
  float_mask = tf.to_float(mask)

    # Hard example mining
  neg_masked_cross_entropy = cross_entropy * (1 - float_mask)


  value1, _ = tf.math.top_k(neg_masked_cross_entropy, k=4096)
  kth1 = tf.reduce_min(value1, 1, keepdims=True)
  mask1 = tf.to_float(tf.less(neg_masked_cross_entropy, kth1))

  value2, _ = tf.math.top_k(tf.multiply(neg_masked_cross_entropy, mask1), k=4096)
  kth2 = tf.reduce_min(value2, 1, keepdims=True)
  mask2 = tf.to_float(tf.less(neg_masked_cross_entropy, kth2))

  value3, _ = tf.math.top_k(tf.multiply(neg_masked_cross_entropy, mask2), k=540)

  value = tf.concat([value1, value2, value3], axis=1)

  num_neg_boxes = tf.minimum(
          tf.to_int32(num_matched_boxes) * ssd_constants.NEGS_PER_POSITIVE, 8731)
  large_neg_ce = tf.batch_gather(value, num_neg_boxes[:, tf.newaxis])
  top_k_neg_mask = tf.to_float(tf.greater_equal(neg_masked_cross_entropy, large_neg_ce))



  class_loss = tf.reduce_sum(
        tf.multiply(cross_entropy, float_mask + top_k_neg_mask), axis=1)


    # TODO(taylorrobie): Confirm that normalizing by the number of boxes matches
    # reference
  return tf.reduce_mean(class_loss / num_matched_boxes)
//  [python]
def _classification_loss(pred_labels, gt_labels, num_matched_boxes):
  """Computes the classification loss.

  Computes the classification loss with hard negative mining.
  Args:
    pred_labels: a dict from index to tensor of predicted class. The shape
      of the tensor is [batch_size, num_anchors, num_classes].
    gt_labels: a list of tensor that represents the classification groundtruth
      targets. The shape is [batch_size, num_anchors, 1].
    num_matched_boxes: the number of anchors that are matched to a groundtruth
      targets. This is used as the loss normalizater.
  Returns:
    box_loss: a float32 representing total box regression loss.
  """
  keys = sorted(pred_labels.keys())
  batch_size = gt_labels[0].shape[0]
  cross_entropy = []
  for i, k in enumerate(keys):
    gt_label = gt_labels[i]
    pred_label = tf.reshape(
        pred_labels[k],
        gt_label.get_shape().as_list() + [ssd_constants.NUM_CLASSES])
    cross_entropy.append(
        tf.reshape(
            _softmax_cross_entropy(pred_label, gt_label), [batch_size, -1]))


  # Put the rest of the loss computation on one device to avoid excessive
  # communication inside topk_mask with spatial partition
  #with tf.device(tf.contrib.tpu.core(0)):
  cross_entropy = tf.concat(cross_entropy, 1)
  gt_label = tf.concat([tf.reshape(l, [batch_size, -1]) for l in gt_labels],
                         1)
  mask = tf.greater(gt_label, 0)
  float_mask = tf.cast(mask, tf.float32)

    # Hard example mining
  neg_masked_cross_entropy = cross_entropy * (1 - float_mask)


  value1, _ = tf.math.top_k(neg_masked_cross_entropy, k=4096)
  kth1 = tf.reduce_min(value1, 1, keepdims=True)
  mask1 = tf.cast(tf.less(neg_masked_cross_entropy, kth1), tf.float32)

  value2, _ = tf.math.top_k(tf.multiply(neg_masked_cross_entropy, mask1), k=4096)
  kth2 = tf.reduce_min(value2, 1, keepdims=True)
  mask2 = tf.cast(tf.less(neg_masked_cross_entropy, kth2), tf.float32)

  value3, _ = tf.math.top_k(tf.multiply(neg_masked_cross_entropy, mask2), k=540)

  value = tf.concat([value1, value2, value3], axis=1)

  num_neg_boxes = tf.minimum(
          tf.to_int32(num_matched_boxes) * ssd_constants.NEGS_PER_POSITIVE, 8731)
  large_neg_ce = tf.batch_gather(value, num_neg_boxes[:, tf.newaxis])
  top_k_neg_mask = tf.cast(tf.greater_equal(neg_masked_cross_entropy, large_neg_ce), tf.float32)



  class_loss = tf.reduce_sum(
        tf.multiply(cross_entropy, float_mask + top_k_neg_mask), axis=1)


    # TODO(taylorrobie): Confirm that normalizing by the number of boxes matches
    # reference
  return tf.reduce_mean(class_loss / num_matched_boxes)
//  [python]
def learning_rate_schedule(params, global_step):
  """Handles learning rate scaling, linear warmup, and learning rate decay.

  Args:
    params: A dictionary that defines hyperparameters of model.
    global_step: A tensor representing current global step.

  Returns:
    A tensor representing current learning rate.
  """
  base_learning_rate = params['base_learning_rate']
  lr_warmup_step = params['lr_warmup_step']
  first_lr_drop_step = params['first_lr_drop_step']
  second_lr_drop_step = params['second_lr_drop_step']
  batch_size = (params['batch_size'] * params['num_shards'] if params['use_tpu']
                else params['batch_size'])
  scaling_factor = batch_size / ssd_constants.DEFAULT_BATCH_SIZE
  mlp_log.mlperf_print('opt_learning_rate_warmup_factor', scaling_factor)
  mlp_log.mlperf_print('opt_learning_rate_warmup_steps', lr_warmup_step)
  adjusted_learning_rate = base_learning_rate * scaling_factor
  learning_rate = (t.cast(global_step, dtype=t.float32) /
                   lr_warmup_step) * adjusted_learning_rate
  lr_schedule = [[1.0, lr_warmup_step], [0.1, first_lr_drop_step],
                 [0.01, second_lr_drop_step]]
  for mult, start_global_step in lr_schedule:
    learning_rate = tf.where(global_step < start_global_step, learning_rate,
                             adjusted_learning_rate * mult)
  return learning_rate
//  [python]
def learning_rate_schedule(params, global_step):
  """Handles learning rate scaling, linear warmup, and learning rate decay.

  Args:
    params: A dictionary that defines hyperparameters of model.
    global_step: A tensor representing current global step.

  Returns:
    A tensor representing current learning rate.
  """
  base_learning_rate = params['base_learning_rate']
  lr_warmup_step = params['lr_warmup_step']
  first_lr_drop_step = params['first_lr_drop_step']
  second_lr_drop_step = params['second_lr_drop_step']
  batch_size = (params['batch_size'] * params['num_shards'] if params['use_tpu']
                else params['batch_size'])
  scaling_factor = batch_size / ssd_constants.DEFAULT_BATCH_SIZE
  mlp_log.mlperf_print('opt_learning_rate_warmup_factor', scaling_factor)
  mlp_log.mlperf_print('opt_learning_rate_warmup_steps', lr_warmup_step)
  adjusted_learning_rate = base_learning_rate * scaling_factor
  learning_rate = (tf.cast(global_step, dtype=tf.float32) /
                   lr_warmup_step) * adjusted_learning_rate
  lr_schedule = [[1.0, lr_warmup_step], [0.1, first_lr_drop_step],
                 [0.01, second_lr_drop_step]]
  for mult, start_global_step in lr_schedule:
    learning_rate = tf.where(global_step < start_global_step, learning_rate,
                             adjusted_learning_rate * mult)
  return learning_rate