jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def relative_positional_encoding(self, qlen, klen, bsz=None, dtype=None):
        """create relative positional encoding."""
        freq_seq = tf.range(0, self.d_model, 2.0)
        if dtype is not None and dtype != tf.float32:
            freq_seq = t.cast(freq_seq, dtype=dtype)
        inv_freq = 1 / (10000 ** (freq_seq / self.d_model))

        if self.attn_type == "bi":
            # beg, end = klen - 1, -qlen
            beg, end = klen, -qlen
        elif self.attn_type == "uni":
            # beg, end = klen - 1, -1
            beg, end = klen, -1
        else:
            raise ValueError("Unknown `attn_type` {}.".format(self.attn_type))

        if self.bi_data:
            fwd_pos_seq = tf.range(beg, end, -1.0)
            bwd_pos_seq = tf.range(-beg, -end, 1.0)

            if dtype is not None and dtype != tf.float32:
                fwd_pos_seq = t.cast(fwd_pos_seq, dtype=dtype)
                bwd_pos_seq = t.cast(bwd_pos_seq, dtype=dtype)

            if self.clamp_len > 0:
                fwd_pos_seq = tf.clip_by_value(fwd_pos_seq, -self.clamp_len, self.clamp_len)
                bwd_pos_seq = tf.clip_by_value(bwd_pos_seq, -self.clamp_len, self.clamp_len)

            if bsz is not None:
                # With bi_data, the batch size should be divisible by 2.
                assert bsz % 2 == 0
                fwd_pos_emb = self.positional_embedding(fwd_pos_seq, inv_freq, bsz // 2)
                bwd_pos_emb = self.positional_embedding(bwd_pos_seq, inv_freq, bsz // 2)
            else:
                fwd_pos_emb = self.positional_embedding(fwd_pos_seq, inv_freq)
                bwd_pos_emb = self.positional_embedding(bwd_pos_seq, inv_freq)

            pos_emb = tf.concat([fwd_pos_emb, bwd_pos_emb], axis=1)
        else:
            fwd_pos_seq = tf.range(beg, end, -1.0)
            if dtype is not None and dtype != tf.float32:
                fwd_pos_seq = t.cast(fwd_pos_seq, dtype=dtype)
            if self.clamp_len > 0:
                fwd_pos_seq = tf.clip_by_value(fwd_pos_seq, -self.clamp_len, self.clamp_len)
            pos_emb = self.positional_embedding(fwd_pos_seq, inv_freq, bsz)

        return pos_emb
//  [python]
def relative_positional_encoding(self, qlen, klen, bsz=None, dtype=None):
        """create relative positional encoding."""
        freq_seq = tf.range(0, self.d_model, 2.0)
        if dtype is not None and dtype != tf.float32:
            freq_seq = tf.cast(freq_seq, dtype=dtype)
        inv_freq = 1 / (10000 ** (freq_seq / self.d_model))

        if self.attn_type == "bi":
            # beg, end = klen - 1, -qlen
            beg, end = klen, -qlen
        elif self.attn_type == "uni":
            # beg, end = klen - 1, -1
            beg, end = klen, -1
        else:
            raise ValueError("Unknown `attn_type` {}.".format(self.attn_type))

        if self.bi_data:
            fwd_pos_seq = tf.range(beg, end, -1.0)
            bwd_pos_seq = tf.range(-beg, -end, 1.0)

            if dtype is not None and dtype != tf.float32:
                fwd_pos_seq = tf.cast(fwd_pos_seq, dtype=dtype)
                bwd_pos_seq = tf.cast(bwd_pos_seq, dtype=dtype)

            if self.clamp_len > 0:
                fwd_pos_seq = tf.clip_by_value(fwd_pos_seq, -self.clamp_len, self.clamp_len)
                bwd_pos_seq = tf.clip_by_value(bwd_pos_seq, -self.clamp_len, self.clamp_len)

            if bsz is not None:
                # With bi_data, the batch size should be divisible by 2.
                assert bsz % 2 == 0
                fwd_pos_emb = self.positional_embedding(fwd_pos_seq, inv_freq, bsz // 2)
                bwd_pos_emb = self.positional_embedding(bwd_pos_seq, inv_freq, bsz // 2)
            else:
                fwd_pos_emb = self.positional_embedding(fwd_pos_seq, inv_freq)
                bwd_pos_emb = self.positional_embedding(bwd_pos_seq, inv_freq)

            pos_emb = tf.concat([fwd_pos_emb, bwd_pos_emb], axis=1)
        else:
            fwd_pos_seq = tf.range(beg, end, -1.0)
            if dtype is not None and dtype != tf.float32:
                fwd_pos_seq = tf.cast(fwd_pos_seq, dtype=dtype)
            if self.clamp_len > 0:
                fwd_pos_seq = tf.clip_by_value(fwd_pos_seq, -self.clamp_len, self.clamp_len)
            pos_emb = self.positional_embedding(fwd_pos_seq, inv_freq, bsz)

        return pos_emb
//  [python]
def scale_boxes(self, patch, ishape, boxes, x_offset, y_offset):
    """Scale and translate the boxes for each image prior to patching."""
    x_offset = t.cast(x_offset, boxes.dtype)
    y_offset = t.cast(y_offset, boxes.dtype)
    pshape = t.cast(t.shape(patch), boxes.dtype)
    ishape = t.cast(ishape, boxes.dtype)
    y_offset = ishape[0] * y_offset
    x_offset = ishape[1] * x_offset

    boxes = box_ops.denormalize_boxes(boxes, pshape[:2])
    boxes = boxes + t.cast(
        [y_offset, x_offset, y_offset, x_offset], boxes.dtype
    )
    boxes = box_ops.normalize_boxes(boxes, ishape[:2])
    return boxes
//  [python]
def scale_boxes(self, patch, ishape, boxes, x_offset, y_offset):
    """Scale and translate the boxes for each image prior to patching."""
    x_offset = tf.cast(x_offset, boxes.dtype)
    y_offset = tf.cast(y_offset, boxes.dtype)
    pshape = tf.cast(tf.shape(patch), boxes.dtype)
    ishape = tf.cast(ishape, boxes.dtype)
    y_offset = ishape[0] * y_offset
    x_offset = ishape[1] * x_offset

    boxes = box_ops.denormalize_boxes(boxes, pshape[:2])
    boxes = boxes + tf.cast(
        [y_offset, x_offset, y_offset, x_offset], boxes.dtype
    )
    boxes = box_ops.normalize_boxes(boxes, ishape[:2])
    return boxes
//  [python]
def _add_param(self, sample):
    """Add parameters to handle skipped images."""
    if 'is_mosaic' not in sample:
      sample['is_mosaic'] = t.cast(0.0, tf.bool)
    sample['num_detections'] = tf.shape(sample['groundtruth_boxes'])[0]
    return sample
//  [python]
def _add_param(self, sample):
    """Add parameters to handle skipped images."""
    if 'is_mosaic' not in sample:
      sample['is_mosaic'] = tf.cast(0.0, tf.bool)
    sample['num_detections'] = tf.shape(sample['groundtruth_boxes'])[0]
    return sample
//  [python]
def call(
      self,
      inputs: tf.Tensor,
      targets: Optional[tf.Tensor] = None,
      training: bool = None,
      use_teacher_forcing_for_eval: bool = False,
      use_input_as_backbone_features=False,
  ) -> List[Any]:
    if use_input_as_backbone_features:
      features = inputs
    else:
      features = self._backbone(inputs)[self._backbone_endpoint_name]
    mask = tf.ones_like(features)
    batch_size, h, w, num_channels = get_shape(features)
    features = tf.reshape(features, [batch_size, h * w, num_channels])
    features = self._stem_ln(
        self._stem_projection(self._dropout(features, training))
    )

    pos_emb = position_embedding_sine(
        mask[:, :, :, 0], num_pos_features=self._hidden_size
    )
    pos_emb = tf.reshape(pos_emb, [batch_size, -1, self._hidden_size])
    pos_emb = t.cast(pos_emb, features.dtype)

    tokens = None
    inputs = {
        "inputs": features,
        "tokens": targets,
        "pos_emb": pos_emb,
    }
    if training:
      logits = self._transformer(inputs, training=True)
    elif use_teacher_forcing_for_eval:
      logits = self._transformer(inputs, training=False)
    else:
      tokens, logits = self._transformer.infer(
          inputs,
          temperature=self._temperature,
          top_k=self._top_k,
          top_p=self._top_p,
          eos_token=self._eos_token,
      )

    return [tokens, logits]
//  [python]
def call(
      self,
      inputs: tf.Tensor,
      targets: Optional[tf.Tensor] = None,
      training: bool = None,
      use_teacher_forcing_for_eval: bool = False,
      use_input_as_backbone_features=False,
  ) -> List[Any]:
    if use_input_as_backbone_features:
      features = inputs
    else:
      features = self._backbone(inputs)[self._backbone_endpoint_name]
    mask = tf.ones_like(features)
    batch_size, h, w, num_channels = get_shape(features)
    features = tf.reshape(features, [batch_size, h * w, num_channels])
    features = self._stem_ln(
        self._stem_projection(self._dropout(features, training))
    )

    pos_emb = position_embedding_sine(
        mask[:, :, :, 0], num_pos_features=self._hidden_size
    )
    pos_emb = tf.reshape(pos_emb, [batch_size, -1, self._hidden_size])
    pos_emb = tf.cast(pos_emb, features.dtype)

    tokens = None
    inputs = {
        "inputs": features,
        "tokens": targets,
        "pos_emb": pos_emb,
    }
    if training:
      logits = self._transformer(inputs, training=True)
    elif use_teacher_forcing_for_eval:
      logits = self._transformer(inputs, training=False)
    else:
      tokens, logits = self._transformer.infer(
          inputs,
          temperature=self._temperature,
          top_k=self._top_k,
          top_p=self._top_p,
          eos_token=self._eos_token,
      )

    return [tokens, logits]
//  [python]
def position_embedding_sine(
    attention_mask,
    num_pos_features=256,
    temperature=10000.0,
    normalize=True,
    scale=2 * math.pi,
):
  """Sine-based positional embeddings for 2D images.

  Args:
    attention_mask: a `bool` Tensor specifying the size of the input image to
      the Transformer and which elements are padded, of size [batch_size,
      height, width]
    num_pos_features: a `int` specifying the number of positional features,
      should be equal to the hidden size of the Transformer network
    temperature: a `float` specifying the temperature of the positional
      embedding. Any type that is converted to a `float` can also be accepted.
    normalize: a `bool` determining whether the positional embeddings should be
      normalized between [0, scale] before application of the sine and cos
      functions.
    scale: a `float` if normalize is True specifying the scale embeddings before
      application of the embedding function.

  Returns:
    embeddings: a `float` tensor of the same shape as input_tensor specifying
      the positional embeddings based on sine features.
  """
  if num_pos_features % 2 != 0:
    raise ValueError(
        "Number of embedding features (num_pos_features) must be even when "
        "column and row embeddings are concatenated."
    )
  num_pos_features = num_pos_features // 2

  # Produce row and column embeddings based on total size of the image
  # <tf.float>[batch_size, height, width]
  attention_mask = t.cast(attention_mask, t.float32)
  row_embedding = tf.cumsum(attention_mask, 1)
  col_embedding = tf.cumsum(attention_mask, 2)

  if normalize:
    eps = 1e-6
    row_embedding = row_embedding / (row_embedding[:, -1:, :] + eps) * scale
    col_embedding = col_embedding / (col_embedding[:, :, -1:] + eps) * scale

  dim_t = tf.range(num_pos_features, dtype=row_embedding.dtype)
  dim_t = tf.pow(temperature, 2 * (dim_t // 2) / num_pos_features)

  # Creates positional embeddings for each row and column position
  # <tf.float>[batch_size, height, width, num_pos_features]
  pos_row = tf.expand_dims(row_embedding, -1) / dim_t
  pos_col = tf.expand_dims(col_embedding, -1) / dim_t
  pos_row = tf.stack(
      [tf.sin(pos_row[:, :, :, 0::2]), tf.cos(pos_row[:, :, :, 1::2])], axis=4
  )
  pos_col = tf.stack(
      [tf.sin(pos_col[:, :, :, 0::2]), tf.cos(pos_col[:, :, :, 1::2])], axis=4
  )

  final_shape = tf_utils.get_shape_list(pos_row)[:3] + [-1]
  pos_row = tf.reshape(pos_row, final_shape)
  pos_col = tf.reshape(pos_col, final_shape)
  output = tf.concat([pos_row, pos_col], -1)

  embeddings = t.cast(output, t.float32)
  return embeddings
//  [python]
def position_embedding_sine(
    attention_mask,
    num_pos_features=256,
    temperature=10000.0,
    normalize=True,
    scale=2 * math.pi,
):
  """Sine-based positional embeddings for 2D images.

  Args:
    attention_mask: a `bool` Tensor specifying the size of the input image to
      the Transformer and which elements are padded, of size [batch_size,
      height, width]
    num_pos_features: a `int` specifying the number of positional features,
      should be equal to the hidden size of the Transformer network
    temperature: a `float` specifying the temperature of the positional
      embedding. Any type that is converted to a `float` can also be accepted.
    normalize: a `bool` determining whether the positional embeddings should be
      normalized between [0, scale] before application of the sine and cos
      functions.
    scale: a `float` if normalize is True specifying the scale embeddings before
      application of the embedding function.

  Returns:
    embeddings: a `float` tensor of the same shape as input_tensor specifying
      the positional embeddings based on sine features.
  """
  if num_pos_features % 2 != 0:
    raise ValueError(
        "Number of embedding features (num_pos_features) must be even when "
        "column and row embeddings are concatenated."
    )
  num_pos_features = num_pos_features // 2

  # Produce row and column embeddings based on total size of the image
  # <tf.float>[batch_size, height, width]
  attention_mask = tf.cast(attention_mask, tf.float32)
  row_embedding = tf.cumsum(attention_mask, 1)
  col_embedding = tf.cumsum(attention_mask, 2)

  if normalize:
    eps = 1e-6
    row_embedding = row_embedding / (row_embedding[:, -1:, :] + eps) * scale
    col_embedding = col_embedding / (col_embedding[:, :, -1:] + eps) * scale

  dim_t = tf.range(num_pos_features, dtype=row_embedding.dtype)
  dim_t = tf.pow(temperature, 2 * (dim_t // 2) / num_pos_features)

  # Creates positional embeddings for each row and column position
  # <tf.float>[batch_size, height, width, num_pos_features]
  pos_row = tf.expand_dims(row_embedding, -1) / dim_t
  pos_col = tf.expand_dims(col_embedding, -1) / dim_t
  pos_row = tf.stack(
      [tf.sin(pos_row[:, :, :, 0::2]), tf.cos(pos_row[:, :, :, 1::2])], axis=4
  )
  pos_col = tf.stack(
      [tf.sin(pos_col[:, :, :, 0::2]), tf.cos(pos_col[:, :, :, 1::2])], axis=4
  )

  final_shape = tf_utils.get_shape_list(pos_row)[:3] + [-1]
  pos_row = tf.reshape(pos_row, final_shape)
  pos_col = tf.reshape(pos_col, final_shape)
  output = tf.concat([pos_row, pos_col], -1)

  embeddings = tf.cast(output, tf.float32)
  return embeddings
//  [python]
def call(self, inputs: tf.Tensor, training: bool = None):
    sources = inputs["inputs"]
    targets = inputs["tokens"]
    mem_pos_embed = inputs["pos_emb"]

    sources = sources + mem_pos_embed
    if self._encoder is not None:
      encoded = self._encoder(sources, None, training=training, ret_list=False)
    else:
      encoded = sources
    encoded = self._output_ln_enc(encoded)

    encoded = self._proj_ln(self._proj(encoded))
    encoded = encoded + mem_pos_embed
    encoded = self._proj_mlp(encoded, training=training)

    seq_len = tf.shape(targets)[1]
    seq_pos_emb = tf.expand_dims(self.seq_pos_emb[:seq_len], 0)
    inp_embedding = outp_embedding = self.token_embedding
    target_emb = tf.gather(inp_embedding, targets) + seq_pos_emb

    self_attention_mask = 1.0 - get_ar_mask(seq_len, target_emb.dtype)

    decoded, _ = self._decoder(
        target_emb, encoded, None, self_attention_mask, None, training
    )
    decoded = self._output_ln_dec(decoded)

    decoded = tf.to_float(decoded)
    outp_embedding = tf.to_float(outp_embedding)

    logits = tf.matmul(decoded, outp_embedding, transpose_b=True)
    if self._output_bias:
      logits = tf.nn.bias_add(logits, self.outp_bias)

    return logits
//  [python]
def call(self, inputs: tf.Tensor, training: bool = None):
    sources = inputs["inputs"]
    targets = inputs["tokens"]
    mem_pos_embed = inputs["pos_emb"]

    sources = sources + mem_pos_embed
    if self._encoder is not None:
      encoded = self._encoder(sources, None, training=training, ret_list=False)
    else:
      encoded = sources
    encoded = self._output_ln_enc(encoded)

    encoded = self._proj_ln(self._proj(encoded))
    encoded = encoded + mem_pos_embed
    encoded = self._proj_mlp(encoded, training=training)

    seq_len = tf.shape(targets)[1]
    seq_pos_emb = tf.expand_dims(self.seq_pos_emb[:seq_len], 0)
    inp_embedding = outp_embedding = self.token_embedding
    target_emb = tf.gather(inp_embedding, targets) + seq_pos_emb

    self_attention_mask = 1.0 - get_ar_mask(seq_len, target_emb.dtype)

    decoded, _ = self._decoder(
        target_emb, encoded, None, self_attention_mask, None, training
    )
    decoded = self._output_ln_dec(decoded)

    decoded = tf.cast(decoded, seq_pos_emb.dtype)
    outp_embedding = tf.cast(outp_embedding, seq_pos_emb.dtype)

    logits = tf.matmul(decoded, outp_embedding, transpose_b=True)
    if self._output_bias:
      logits = tf.nn.bias_add(logits, self.outp_bias)

    return logits
//  [python]
def _update_patched_sample(
      self, sample, image, boxes, classes, is_crowds, areas, mosaic_center
  ):
    """Returns a shallow copy of sample with updated values."""
    boxes = tf.concat(boxes, axis=0)
    classes = tf.concat(classes, axis=0)
    is_crowds = tf.concat(is_crowds, axis=0)
    areas = tf.concat(areas, axis=0)

    if self._mosaic_crop_mode is not None:
      image, boxes, classes, is_crowds, areas, _ = self._mosaic_crop_image(
          image, boxes, classes, is_crowds, areas, mosaic_center
      )

    height, width = preprocessing_ops.get_image_shape(image)
    # Shallow copy of dict is needed to keep this method functional and
    # AutoGraph happy.
    sample = sample.copy()
    sample['image'] = t.cast(image, t.uint8)
    sample['groundtruth_boxes'] = boxes
    sample['groundtruth_area'] = areas
    sample['groundtruth_classes'] = t.cast(
        classes, sample['groundtruth_classes'].dtype
    )
    sample['groundtruth_is_crowd'] = t.cast(is_crowds, t.bool)
    sample['width'] = t.cast(width, sample['width'].dtype)
    sample['height'] = t.cast(height, sample['height'].dtype)
    sample['num_detections'] = tf.shape(sample['groundtruth_boxes'])[1]
    sample['is_mosaic'] = t.cast(1.0, t.bool)

    del sample['shiftx']
    del sample['shifty']
    del sample['crop_points']

    return sample
//  [python]
def _update_patched_sample(
      self, sample, image, boxes, classes, is_crowds, areas, mosaic_center
  ):
    """Returns a shallow copy of sample with updated values."""
    boxes = tf.concat(boxes, axis=0)
    classes = tf.concat(classes, axis=0)
    is_crowds = tf.concat(is_crowds, axis=0)
    areas = tf.concat(areas, axis=0)

    if self._mosaic_crop_mode is not None:
      image, boxes, classes, is_crowds, areas, _ = self._mosaic_crop_image(
          image, boxes, classes, is_crowds, areas, mosaic_center
      )

    height, width = preprocessing_ops.get_image_shape(image)
    # Shallow copy of dict is needed to keep this method functional and
    # AutoGraph happy.
    sample = sample.copy()
    sample['image'] = tf.cast(image, tf.uint8)
    sample['groundtruth_boxes'] = boxes
    sample['groundtruth_area'] = areas
    sample['groundtruth_classes'] = tf.cast(
        classes, sample['groundtruth_classes'].dtype
    )
    sample['groundtruth_is_crowd'] = tf.cast(is_crowds, tf.bool)
    sample['width'] = tf.cast(width, sample['width'].dtype)
    sample['height'] = tf.cast(height, sample['height'].dtype)
    sample['num_detections'] = tf.shape(sample['groundtruth_boxes'])[1]
    sample['is_mosaic'] = tf.cast(1.0, tf.bool)

    del sample['shiftx']
    del sample['shifty']
    del sample['crop_points']

    return sample
//  [python]
def _mixup(self, one, two):
    """Blend together 2 images for the mixup data augmentation."""
    if self._mixup_frequency >= 1.0:
      domo = 1.0
    else:
      domo = preprocessing_ops.random_uniform_strong(
          0.0, 1.0, dtype=tf.float32, seed=self._seed)
      noop = one.copy()

    if domo >= (1 - self._mixup_frequency):
      sample = one
      otype = one['image'].dtype

      r = self._beta(8.0, 8.0)
      sample['image'] = (
          r * t.cast(one['image'], tf.float32) +
          (1 - r) * t.cast(two['image'], tf.float32))

      sample['image'] = t.cast(sample['image'], otype)
      sample['groundtruth_boxes'] = tf.concat(
          [one['groundtruth_boxes'], two['groundtruth_boxes']], axis=0)
      sample['groundtruth_classes'] = tf.concat(
          [one['groundtruth_classes'], two['groundtruth_classes']], axis=0)
      sample['groundtruth_is_crowd'] = tf.concat(
          [one['groundtruth_is_crowd'], two['groundtruth_is_crowd']], axis=0)
      sample['groundtruth_area'] = tf.concat(
          [one['groundtruth_area'], two['groundtruth_area']], axis=0)
      return sample
    else:
      return self._add_param(noop)
//  [python]
def _mixup(self, one, two):
    """Blend together 2 images for the mixup data augmentation."""
    if self._mixup_frequency >= 1.0:
      domo = 1.0
    else:
      domo = preprocessing_ops.random_uniform_strong(
          0.0, 1.0, dtype=tf.float32, seed=self._seed)
      noop = one.copy()

    if domo >= (1 - self._mixup_frequency):
      sample = one
      otype = one['image'].dtype

      r = self._beta(8.0, 8.0)
      sample['image'] = (
          r * tf.cast(one['image'], tf.float32) +
          (1 - r) * tf.cast(two['image'], tf.float32))

      sample['image'] = tf.cast(sample['image'], otype)
      sample['groundtruth_boxes'] = tf.concat(
          [one['groundtruth_boxes'], two['groundtruth_boxes']], axis=0)
      sample['groundtruth_classes'] = tf.concat(
          [one['groundtruth_classes'], two['groundtruth_classes']], axis=0)
      sample['groundtruth_is_crowd'] = tf.concat(
          [one['groundtruth_is_crowd'], two['groundtruth_is_crowd']], axis=0)
      sample['groundtruth_area'] = tf.concat(
          [one['groundtruth_area'], two['groundtruth_area']], axis=0)
      return sample
    else:
      return self._add_param(noop)
//  [python]
def _mosaic_crop_image(
      self, image, boxes, classes, is_crowd, area, mosaic_center):
    """Process a patched image in preperation for final output."""
    if self._mosaic_crop_mode != 'crop':
      shape = t.cast(preprocessing_ops.get_image_shape(image), t.float32)
      center = shape * mosaic_center

      # shift the center of the image by applying a translation to the whole
      # image
      ch = t.math.round(
          preprocessing_ops.random_uniform_strong(
              -center[0], center[0], seed=self._seed))
      cw = t.math.round(
          preprocessing_ops.random_uniform_strong(
              -center[1], center[1], seed=self._seed))

      # clip the boxes to fit within the image
      image = augment.translate(
          image, [cw, ch], fill_value=self._pad_value, fill_mode='constant'
      )
      boxes = box_ops.denormalize_boxes(boxes, shape[:2])
      boxes = boxes + t.cast([ch, cw, ch, cw], boxes.dtype)
      boxes = box_ops.clip_boxes(boxes, shape[:2])
      inds = box_ops.get_non_empty_box_indices(boxes)

      boxes = box_ops.normalize_boxes(boxes, shape[:2])
      boxes, classes, is_crowd, area = self._select_ind(inds, boxes, classes,  # pylint:disable=unbalanced-tuple-unpacking
                                                        is_crowd, area)

    # warp and scale the fully stitched sample
    image, _, affine = preprocessing_ops.affine_warp_image(
        image, [self._output_size[0], self._output_size[1]],
        scale_min=self._aug_scale_min,
        scale_max=self._aug_scale_max,
        translate=self._aug_rand_translate,
        degrees=self._aug_rand_angle,
        perspective=self._aug_rand_perspective,
        random_pad=self._random_pad,
        seed=self._seed)
    height, width = self._output_size[0], self._output_size[1]
    image = t.image.resize(image, (height, width))

    # clip and clean boxes
    boxes, inds = preprocessing_ops.transform_and_clip_boxes(
        boxes,
        None,
        affine=affine,
        area_thresh=self._area_thresh,
        seed=self._seed)
    classes, is_crowd, area = self._select_ind(inds, classes, is_crowd, area)  # pylint:disable=unbalanced-tuple-unpacking
    return image, boxes, classes, is_crowd, area, area
//  [python]
def _mosaic_crop_image(
      self, image, boxes, classes, is_crowd, area, mosaic_center):
    """Process a patched image in preperation for final output."""
    if self._mosaic_crop_mode != 'crop':
      shape = tf.cast(preprocessing_ops.get_image_shape(image), tf.float32)
      center = shape * mosaic_center

      # shift the center of the image by applying a translation to the whole
      # image
      ch = tf.math.round(
          preprocessing_ops.random_uniform_strong(
              -center[0], center[0], seed=self._seed))
      cw = tf.math.round(
          preprocessing_ops.random_uniform_strong(
              -center[1], center[1], seed=self._seed))

      # clip the boxes to fit within the image
      image = augment.translate(
          image, [cw, ch], fill_value=self._pad_value, fill_mode='constant'
      )
      boxes = box_ops.denormalize_boxes(boxes, shape[:2])
      boxes = boxes + tf.cast([ch, cw, ch, cw], boxes.dtype)
      boxes = box_ops.clip_boxes(boxes, shape[:2])
      inds = box_ops.get_non_empty_box_indices(boxes)

      boxes = box_ops.normalize_boxes(boxes, shape[:2])
      boxes, classes, is_crowd, area = self._select_ind(inds, boxes, classes,  # pylint:disable=unbalanced-tuple-unpacking
                                                        is_crowd, area)

    # warp and scale the fully stitched sample
    image, _, affine = preprocessing_ops.affine_warp_image(
        image, [self._output_size[0], self._output_size[1]],
        scale_min=self._aug_scale_min,
        scale_max=self._aug_scale_max,
        translate=self._aug_rand_translate,
        degrees=self._aug_rand_angle,
        perspective=self._aug_rand_perspective,
        random_pad=self._random_pad,
        seed=self._seed)
    height, width = self._output_size[0], self._output_size[1]
    image = tf.image.resize(image, (height, width))

    # clip and clean boxes
    boxes, inds = preprocessing_ops.transform_and_clip_boxes(
        boxes,
        None,
        affine=affine,
        area_thresh=self._area_thresh,
        seed=self._seed)
    classes, is_crowd, area = self._select_ind(inds, classes, is_crowd, area)  # pylint:disable=unbalanced-tuple-unpacking
    return image, boxes, classes, is_crowd, area, area
//  [python]
def _get_proximity_to_anchor(self, action, anchor_point, dist):
    delta = t.cast(action, tf.float32) - anchor_point
    logging.info('Test q_func \nanchor: %s, dist: %s', anchor_point, dist)
    if dist == DIST_L1:
      val = -tf.abs(delta)
    elif dist == DIST_BIMODAL:
      val = tfp.distributions.Normal(0.0, VAR_ANCHOR_PT).prob(
          delta
      ) + tfp.distributions.Normal(0.0, VAR_LOCAL_MAX).prob(
          action - LOCAL_MAX_BIMODAL
      )
    else:
      raise ValueError('Unknown distance option for Q-function definition.')
    return val
//  [python]
def _get_proximity_to_anchor(self, action, anchor_point, dist):
    delta = tf.cast(action, tf.float32) - anchor_point
    logging.info('Test q_func \nanchor: %s, dist: %s', anchor_point, dist)
    if dist == DIST_L1:
      val = -tf.abs(delta)
    elif dist == DIST_BIMODAL:
      val = tfp.distributions.Normal(0.0, VAR_ANCHOR_PT).prob(
          delta
      ) + tfp.distributions.Normal(0.0, VAR_LOCAL_MAX).prob(
          action - LOCAL_MAX_BIMODAL
      )
    else:
      raise ValueError('Unknown distance option for Q-function definition.')
    return val
//  [python]
def call(self, inputs, step_type=(), network_state=()):
    _, action_nested = inputs

    flat_action = tf.nest.flatten(action_nested)

    val = self._get_proximity_to_anchor(
        flat_action[self.idx_continuous], self._anchor_point, self._dist
    )
    res = tf.reduce_min(val, axis=-1)

    # If there are categorical actions, then make the Q value very high
    # when they have high returns in categorical_action_returns.
    # The expectation is that CEM will find this the best
    # possible action since this Q value dominates the ones coming from
    # proximity to the given anchor point.
    if self._sampler_type == 'continuous_and_one_hot':
      cat_actions_dist_1 = tf.abs(flat_action[1 - self.idx_continuous] - 1)
      cat_returns = tf.constant(self._categorical_action_returns)
      overall_cat_actions_vals = -tf.multiply(
          cat_returns, t.cast(cat_actions_dist_1, tf.float32)
      )
      res += tf.reduce_min(overall_cat_actions_vals, axis=1)
    elif self._sampler_type == 'hybrid':
      val_discrete = self._get_proximity_to_anchor(
          flat_action[1 - self.idx_continuous], self._anchor_point, self._dist
      )
      res += tf.reduce_min(val_discrete, axis=-1)

    return res, ()
//  [python]
def call(self, inputs, step_type=(), network_state=()):
    _, action_nested = inputs

    flat_action = tf.nest.flatten(action_nested)

    val = self._get_proximity_to_anchor(
        flat_action[self.idx_continuous], self._anchor_point, self._dist
    )
    res = tf.reduce_min(val, axis=-1)

    # If there are categorical actions, then make the Q value very high
    # when they have high returns in categorical_action_returns.
    # The expectation is that CEM will find this the best
    # possible action since this Q value dominates the ones coming from
    # proximity to the given anchor point.
    if self._sampler_type == 'continuous_and_one_hot':
      cat_actions_dist_1 = tf.abs(flat_action[1 - self.idx_continuous] - 1)
      cat_returns = tf.constant(self._categorical_action_returns)
      overall_cat_actions_vals = -tf.multiply(
          cat_returns, tf.cast(cat_actions_dist_1, tf.float32)
      )
      res += tf.reduce_min(overall_cat_actions_vals, axis=1)
    elif self._sampler_type == 'hybrid':
      val_discrete = self._get_proximity_to_anchor(
          flat_action[1 - self.idx_continuous], self._anchor_point, self._dist
      )
      res += tf.reduce_min(val_discrete, axis=-1)

    return res, ()
//  [python]
def accuracy_op(predictions, targets):
    """ accuracy_op.

    An op that calculates mean accuracy, assuming predictiosn are targets
    are both one-hot encoded.

    Examples:
        
        input_data = placeholder(shape=[None, 784])
        y_pred = my_network(input_data) # Apply some ops
        y_true = placeholder(shape=[None, 10]) # Labels
        acc_op = accuracy_op(y_pred, y_true)

        # Calculate accuracy by feeding data X and labels Y
        accuracy = sess.run(acc_op, feed_dict={input_data: X, y_true: Y})
        

    Arguments:
        predictions: `Tensor`.
        targets: `Tensor`.

    Returns:
        `Float`. The mean accuracy.

    """
    if not isinstance(targets, tf.Tensor):
        raise ValueError("mean_accuracy 'input' argument only accepts type "
                         "Tensor, '" + str(type(input)) + "' given.")

    with tf.name_scope('Accuracy'):
        correct_pred = tf.equal(tf.argmax(predictions, 1), tf.argmax(targets, 1))
        acc = tf.reduce_mean(tf.to_float(correct_pred))
    return acc
//  [python]
def accuracy_op(predictions, targets):
    """ accuracy_op.

    An op that calculates mean accuracy, assuming predictiosn are targets
    are both one-hot encoded.

    Examples:
        ```python
        input_data = placeholder(shape=[None, 784])
        y_pred = my_network(input_data) # Apply some ops
        y_true = placeholder(shape=[None, 10]) # Labels
        acc_op = accuracy_op(y_pred, y_true)

        # Calculate accuracy by feeding data X and labels Y
        accuracy = sess.run(acc_op, feed_dict={input_data: X, y_true: Y})
        ```

    Arguments:
        predictions: `Tensor`.
        targets: `Tensor`.

    Returns:
        `Float`. The mean accuracy.

    """
    if not isinstance(targets, tf.Tensor):
        raise ValueError("mean_accuracy 'input' argument only accepts type "
                         "Tensor, '" + str(type(input)) + "' given.")

    with tf.name_scope('Accuracy'):
        correct_pred = tf.equal(tf.argmax(predictions, 1), tf.argmax(targets, 1))
        acc = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
    return acc
//  [python]
def _assert_correctness(
      self, res, sampler_type, anchor_point=0.0, categorical_action_returns=None
  ):
    logging.info('anchor = %.2f', anchor_point)
    logging.info('Output best actions (size [batch, action])\n %s', res)

    if sampler_type == 'continuous_and_one_hot':
      num_categorical_actions = len(categorical_action_returns)
      argmax_action = np.argmax(categorical_action_returns)

      res = self.evaluate(res)

      # First ensure all continuous parts are 0 for whole batch.
      self.assertEqual(res['continuous'].all(), 0.0)
      # Ensure the right categorical action is chosen for whole batch.
      expected = [
          1 if i == argmax_action else 0 for i in range(num_categorical_actions)
      ]
      expected = np.stack([expected for i in range(len(res))])
      self.assertAllEqual(expected, res['discrete'][:, 0, :])
    elif sampler_type == 'hybrid':
      # All actions must be close to the anchor point.
      self.assertAllLessEqual(
          tf.abs(tf.cast(res['discrete'], tf.float32) - anchor_point), DELTA
      )
      self.assertAllLessEqual(tf.abs(res['continuous'] - anchor_point), DELTA)
    elif sampler_type == 'continuous':
      # All actions must be close to the anchor point.
      self.assertAllLessEqual(tf.abs(res - anchor_point), DELTA)
//  [python]
def top_k_op(predictions, targets, k=1):
    """ top_k_op.

    An op that calculates top-k mean accuracy.

    Examples:
        
        input_data = placeholder(shape=[None, 784])
        y_pred = my_network(input_data) # Apply some ops
        y_true = placeholder(shape=[None, 10]) # Labels
        top3_op = top_k_op(y_pred, y_true, 3)

        # Calculate Top-3 accuracy by feeding data X and labels Y
        top3_accuracy = sess.run(top3_op, feed_dict={input_data: X, y_true: Y})
        

    Arguments:
        predictions: `Tensor`.
        targets: `Tensor`.
        k: `int`. Number of top elements to look at for computing precision.

    Returns:
        `Float`. The top-k mean accuracy.

    """
    with tf.name_scope('Top_' + str(k)):
        targets = tf.to_int32(targets)
        correct_pred = tf.nn.in_top_k(predictions, tf.argmax(targets, 1), k)
        acc = tf.reduce_mean(tf.to_float(correct_pred))
    return acc
//  [python]
def top_k_op(predictions, targets, k=1):
    """ top_k_op.

    An op that calculates top-k mean accuracy.

    Examples:
        ```python
        input_data = placeholder(shape=[None, 784])
        y_pred = my_network(input_data) # Apply some ops
        y_true = placeholder(shape=[None, 10]) # Labels
        top3_op = top_k_op(y_pred, y_true, 3)

        # Calculate Top-3 accuracy by feeding data X and labels Y
        top3_accuracy = sess.run(top3_op, feed_dict={input_data: X, y_true: Y})
        ```

    Arguments:
        predictions: `Tensor`.
        targets: `Tensor`.
        k: `int`. Number of top elements to look at for computing precision.

    Returns:
        `Float`. The top-k mean accuracy.

    """
    with tf.name_scope('Top_' + str(k)):
        targets = tf.cast(targets, tf.int32)
        correct_pred = tf.nn.in_top_k(predictions, tf.argmax(targets, 1), k)
        acc = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
    return acc
//  [python]
def binary_accuracy_op(predictions, targets):
    """ binary_accuracy_op.

    An op that calculates mean accuracy, assuming predictions are logits, and
    targets are binary encoded (and represented as int32).

    Examples:
        
        input_data = placeholder(shape=[None, 784])
        y_pred = my_network(input_data) # Apply some ops
        y_true = placeholder(shape=[None, 10]) # Labels
        acc_op = binary_accuracy_op(y_pred, y_true)

        # Calculate accuracy by feeding data X and labels Y
        binary_accuracy = sess.run(acc_op, feed_dict={input_data: X, y_true: Y})
        

    Arguments:
        predictions: `Tensor` of `float` type.
        targets: `Tensor` of `float` type.

    Returns:
        `Float`. The mean accuracy.

    """
    if not isinstance(targets, tf.Tensor):
        raise ValueError("mean_accuracy 'input' argument only accepts type "
                         "Tensor, '" + str(type(input)) + "' given.")

    with tf.name_scope('BinaryAccuracy'):
        predictions = tf.to_float(tf.greater(predictions, 0))
        correct_pred = tf.equal(predictions, tf.to_float(targets))
        acc = tf.reduce_mean(tf.to_float(correct_pred))
    return acc
//  [python]
def binary_accuracy_op(predictions, targets):
    """ binary_accuracy_op.

    An op that calculates mean accuracy, assuming predictions are logits, and
    targets are binary encoded (and represented as int32).

    Examples:
        ```python
        input_data = placeholder(shape=[None, 784])
        y_pred = my_network(input_data) # Apply some ops
        y_true = placeholder(shape=[None, 10]) # Labels
        acc_op = binary_accuracy_op(y_pred, y_true)

        # Calculate accuracy by feeding data X and labels Y
        binary_accuracy = sess.run(acc_op, feed_dict={input_data: X, y_true: Y})
        ```

    Arguments:
        predictions: `Tensor` of `float` type.
        targets: `Tensor` of `float` type.

    Returns:
        `Float`. The mean accuracy.

    """
    if not isinstance(targets, tf.Tensor):
        raise ValueError("mean_accuracy 'input' argument only accepts type "
                         "Tensor, '" + str(type(input)) + "' given.")

    with tf.name_scope('BinaryAccuracy'):
        predictions = tf.cast(tf.greater(predictions, 0), tf.float32)
        correct_pred = tf.equal(predictions, tf.cast(targets, tf.float32))
        acc = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
    return acc
//  [python]
def validation_step(self,
                      inputs: Tuple[Any, Any],
                      model: tf_keras.Model,
                      metrics: Optional[List[Any]] = None):
    """Validatation step.

    Args:
      inputs: a dictionary of input tensors.
      model: the keras.Model.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    images, labels = inputs
    outputs = model(
        images,
        anchor_boxes=labels['anchor_boxes'],
        image_shape=labels['image_info'][:, 1, :],
        training=False,
    )

    logs = {self.loss: 0}
    self._update_metrics(labels, outputs, logs)

    if (
        hasattr(self.task_config, 'allow_image_summary')
        and self.task_config.allow_image_summary
    ):
      logs.update(
          {'visualization': (t.cast(images, dtype=t.float32), outputs)}
      )

    return logs
//  [python]
def validation_step(self,
                      inputs: Tuple[Any, Any],
                      model: tf_keras.Model,
                      metrics: Optional[List[Any]] = None):
    """Validatation step.

    Args:
      inputs: a dictionary of input tensors.
      model: the keras.Model.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    images, labels = inputs
    outputs = model(
        images,
        anchor_boxes=labels['anchor_boxes'],
        image_shape=labels['image_info'][:, 1, :],
        training=False,
    )

    logs = {self.loss: 0}
    self._update_metrics(labels, outputs, logs)

    if (
        hasattr(self.task_config, 'allow_image_summary')
        and self.task_config.allow_image_summary
    ):
      logs.update(
          {'visualization': (tf.cast(images, dtype=tf.float32), outputs)}
      )

    return logs
//  [python]
def _normalize(self, input_image, input_mask):
        """ Normalise input image
        Args:
            input_image (tf.image): The input image
            input_mask (int): The image mask

        Returns:
            input_image (tf.image): The normalized input image
            input_mask (int): The new image mask
        """
        input_image = tf.to_float(input_image) / 255.0
        input_mask -= 1
        return input_image, input_mask
//  [python]
def _normalize(self, input_image, input_mask):
        """ Normalise input image
        Args:
            input_image (tf.image): The input image
            input_mask (int): The image mask

        Returns:
            input_image (tf.image): The normalized input image
            input_mask (int): The new image mask
        """
        input_image = tf.cast(input_image, tf.float32) / 255.0
        input_mask -= 1
        return input_image, input_mask
//  [python]
def train_step(self,
                 inputs: Tuple[Any, Any],
                 model: tf_keras.Model,
                 optimizer: tf_keras.optimizers.Optimizer,
                 metrics: Optional[List[Any]] = None):
    """Does forward and backward.

    Args:
      inputs: a dictionary of input tensors.
      model: the model, forward pass definition.
      optimizer: the optimizer for this training step.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    images, labels = inputs
    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
    with tf.GradientTape() as tape:
      model_kwargs = {
          'image_shape': labels['image_info'][:, 1, :],
          'anchor_boxes': labels['anchor_boxes'],
          'gt_boxes': labels['gt_boxes'],
          'gt_classes': labels['gt_classes'],
          'training': True,
      }
      if self.task_config.model.include_mask:
        model_kwargs['gt_masks'] = labels['gt_masks']
        if self.task_config.model.outer_boxes_scale > 1.0:
          model_kwargs['gt_outer_boxes'] = labels['gt_outer_boxes']
      outputs = model(
          images, **model_kwargs)
      outputs = tf.nest.map_structure(
          lambda x: tf.to_float(x), outputs)

      # Computes per-replica loss.
      losses = self.build_losses(
          outputs=outputs, labels=labels, aux_losses=model.losses)
      scaled_loss = losses['total_loss'] / num_replicas

      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
      # scaled for numerical stability.
      if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    tvars = model.trainable_variables
    grads = tape.gradient(scaled_loss, tvars)
    # Scales back gradient when LossScaleOptimizer is used.
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      grads = optimizer.get_unscaled_gradients(grads)
    optimizer.apply_gradients(list(zip(grads, tvars)))

    logs = {self.loss: losses['total_loss']}

    if metrics:
      for m in metrics:
        m.update_state(losses[m.name])

    return logs
//  [python]
def train_step(self,
                 inputs: Tuple[Any, Any],
                 model: tf_keras.Model,
                 optimizer: tf_keras.optimizers.Optimizer,
                 metrics: Optional[List[Any]] = None):
    """Does forward and backward.

    Args:
      inputs: a dictionary of input tensors.
      model: the model, forward pass definition.
      optimizer: the optimizer for this training step.
      metrics: a nested structure of metrics objects.

    Returns:
      A dictionary of logs.
    """
    images, labels = inputs
    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
    with tf.GradientTape() as tape:
      model_kwargs = {
          'image_shape': labels['image_info'][:, 1, :],
          'anchor_boxes': labels['anchor_boxes'],
          'gt_boxes': labels['gt_boxes'],
          'gt_classes': labels['gt_classes'],
          'training': True,
      }
      if self.task_config.model.include_mask:
        model_kwargs['gt_masks'] = labels['gt_masks']
        if self.task_config.model.outer_boxes_scale > 1.0:
          model_kwargs['gt_outer_boxes'] = labels['gt_outer_boxes']
      outputs = model(
          images, **model_kwargs)
      outputs = tf.nest.map_structure(
          lambda x: tf.cast(x, tf.float32), outputs)

      # Computes per-replica loss.
      losses = self.build_losses(
          outputs=outputs, labels=labels, aux_losses=model.losses)
      scaled_loss = losses['total_loss'] / num_replicas

      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
      # scaled for numerical stability.
      if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    tvars = model.trainable_variables
    grads = tape.gradient(scaled_loss, tvars)
    # Scales back gradient when LossScaleOptimizer is used.
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      grads = optimizer.get_unscaled_gradients(grads)
    optimizer.apply_gradients(list(zip(grads, tvars)))

    logs = {self.loss: losses['total_loss']}

    if metrics:
      for m in metrics:
        m.update_state(losses[m.name])

    return logs
//  [python]
def compute_Fl(flow_gt, flow_est, mask):
    # F1 measure
    err = tf.multiply(flow_gt - flow_est, mask)
    err_norm = tf.norm(err, axis=-1)
    
    flow_gt_norm = tf.maximum(tf.norm(flow_gt, axis=-1), 1e-12)
    F1_logic = tf.logical_and(err_norm > 3, tf.divide(err_norm, flow_gt_norm) > 0.05)
    F1_logic = t.cast(tf.logical_and(tf.expand_dims(F1_logic, -1), mask > 0), tf.float32)
    F1 = tf.reduce_sum(F1_logic) / (tf.reduce_sum(mask) + 1e-6)
    return F1
//  [python]
def compute_Fl(flow_gt, flow_est, mask):
    # F1 measure
    err = tf.multiply(flow_gt - flow_est, mask)
    err_norm = tf.norm(err, axis=-1)
    
    flow_gt_norm = tf.maximum(tf.norm(flow_gt, axis=-1), 1e-12)
    F1_logic = tf.logical_and(err_norm > 3, tf.divide(err_norm, flow_gt_norm) > 0.05)
    F1_logic = tf.cast(tf.logical_and(tf.expand_dims(F1_logic, -1), mask > 0), tf.float32)
    F1 = tf.reduce_sum(F1_logic) / (tf.reduce_sum(mask) + 1e-6)
    return F1
//  [python]
def get_learning_rate(train_config, global_step=None):
    """
    Get learning rate from train config.

    TODO: Better config usage.

    Returns:
        learning_rate: TensorFlow variable.

    Raises:
        ValueError: When the method used is not available.
    """
    lr_config = train_config.learning_rate.copy()
    decay_method = lr_config.pop('decay_method', None)

    if not decay_method or decay_method == 'none':
        return lr_config.get('value') or lr_config.get('learning_rate')

    if decay_method not in LEARNING_RATE_DECAY_METHODS:
        raise ValueError('Invalid learning_rate method "{}"'.format(
            decay_method
        ))

    if decay_method == 'piecewise_constant':
        lr_config['x'] = global_step
    else:
        lr_config['global_step'] = global_step

    # boundaries, when used, must be the same type as global_step (int64).
    if 'boundaries' in lr_config:
        lr_config['boundaries'] = [
            t.cast(b, t.int64) for b in lr_config['boundaries']
        ]

    decay_function = LEARNING_RATE_DECAY_METHODS[decay_method]
    learning_rate = decay_function(
        **lr_config
    )

    tf.summary.scalar('losses/learning_rate', learning_rate)

    return learning_rate
//  [python]
def get_learning_rate(train_config, global_step=None):
    """
    Get learning rate from train config.

    TODO: Better config usage.

    Returns:
        learning_rate: TensorFlow variable.

    Raises:
        ValueError: When the method used is not available.
    """
    lr_config = train_config.learning_rate.copy()
    decay_method = lr_config.pop('decay_method', None)

    if not decay_method or decay_method == 'none':
        return lr_config.get('value') or lr_config.get('learning_rate')

    if decay_method not in LEARNING_RATE_DECAY_METHODS:
        raise ValueError('Invalid learning_rate method "{}"'.format(
            decay_method
        ))

    if decay_method == 'piecewise_constant':
        lr_config['x'] = global_step
    else:
        lr_config['global_step'] = global_step

    # boundaries, when used, must be the same type as global_step (int64).
    if 'boundaries' in lr_config:
        lr_config['boundaries'] = [
            tf.cast(b, tf.int64) for b in lr_config['boundaries']
        ]

    decay_function = LEARNING_RATE_DECAY_METHODS[decay_method]
    learning_rate = decay_function(
        **lr_config
    )

    tf.summary.scalar('losses/learning_rate', learning_rate)

    return learning_rate