jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def input_fn_builder(input_file, seq_length, is_training, drop_remainder):
    """Creates an `input_fn` closure to be passed to TPUEstimator."""

    name_to_features = {
        "unique_ids": tf.FixedLenFeature([], tf.int64),
        "input_ids": tf.FixedLenFeature([seq_length], tf.int64),
        "input_mask": tf.FixedLenFeature([seq_length], tf.int64),
        "segment_ids": tf.FixedLenFeature([seq_length], tf.int64),
    }

    if is_training:
        name_to_features["start_positions"] = tf.FixedLenFeature([], tf.int64)
        name_to_features["end_positions"] = tf.FixedLenFeature([], tf.int64)

    def _decode_record(record, name_to_features):
        """Decodes a record to a TensorFlow example."""
        example = tf.parse_single_example(record, name_to_features)

        # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
        # So cast all int64 to int32.
        for name in list(example.keys()):
            t = example[name]
            if t.dtype == tf.int64:
                t = tf.cast(t, tf.int32)
            example[name] = t

        return example

    def input_fn(params):
        """The actual input function."""
        batch_size = params["batch_size"]

        # For training, we want a lot of parallel reading and shuffling.
        # For eval, we want no shuffling and parallel reading doesn't matter.
        d = tf.data.TFRecordDataset(input_file)
        if is_training:
            d = d.repeat()
            d = d.shuffle(buffer_size=100)

        d = d.apply(
            tf.data.experimental.map_and_batch(
                lambda record: _decode_record(record, name_to_features),
                batch_size=batch_size,
                drop_remainder=drop_remainder,
            )
        )

        return d

    return input_fn
//  [python]
def _shift_bbox(bbox,
                image_height,
                image_width,
                pixels,
                shift_horizontal):
  """Shifts the bbox coordinates by pixels.

  Args:
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    image_height: Int, height of the image.
    image_width: Int, width of the image.
    pixels: An int. How many pixels to shift the bbox.
    shift_horizontal: Boolean. If true then shift in X dimension else shift in
      Y dimension.

  Returns:
    A tensor of the same shape as bbox, but now with the shifted coordinates.
  """
  pixels = tf.to_int32(pixels)
  # Convert bbox to integer pixel locations.
  min_y = tf.to_int32(tf.to_float(image_height) * bbox[0])
  min_x = tf.to_int32(tf.to_float(image_width) * bbox[1])
  max_y = tf.to_int32(tf.to_float(image_height) * bbox[2])
  max_x = tf.to_int32(tf.to_float(image_width) * bbox[3])

  if shift_horizontal:
    min_x = tf.maximum(0, min_x - pixels)
    max_x = tf.minimum(image_width, max_x - pixels)
  else:
    min_y = tf.maximum(0, min_y - pixels)
    max_y = tf.minimum(image_height, max_y - pixels)

  # Convert bbox back to floats.
  min_y = tf.to_float(min_y) / tf.to_float(image_height)
  min_x = tf.to_float(min_x) / tf.to_float(image_width)
  max_y = tf.to_float(max_y) / tf.to_float(image_height)
  max_x = tf.to_float(max_x) / tf.to_float(image_width)

  # Clip the bboxes to be sure the fall between [0, 1].
  min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
  min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
  return tf.stack([min_y, min_x, max_y, max_x])
//  [python]
def _shift_bbox(bbox,
                image_height,
                image_width,
                pixels,
                shift_horizontal):
  """Shifts the bbox coordinates by pixels.

  Args:
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    image_height: Int, height of the image.
    image_width: Int, width of the image.
    pixels: An int. How many pixels to shift the bbox.
    shift_horizontal: Boolean. If true then shift in X dimension else shift in
      Y dimension.

  Returns:
    A tensor of the same shape as bbox, but now with the shifted coordinates.
  """
  pixels = tf.cast(pixels, tf.int32)
  # Convert bbox to integer pixel locations.
  min_y = tf.cast(tf.to_float(image_height) * bbox[0], tf.int32)
  min_x = tf.cast(tf.to_float(image_width) * bbox[1], tf.int32)
  max_y = tf.cast(tf.to_float(image_height) * bbox[2], tf.int32)
  max_x = tf.cast(tf.to_float(image_width) * bbox[3], tf.int32)

  if shift_horizontal:
    min_x = tf.maximum(0, min_x - pixels)
    max_x = tf.minimum(image_width, max_x - pixels)
  else:
    min_y = tf.maximum(0, min_y - pixels)
    max_y = tf.minimum(image_height, max_y - pixels)

  # Convert bbox back to floats.
  min_y = tf.to_float(min_y) / tf.to_float(image_height)
  min_x = tf.to_float(min_x) / tf.to_float(image_width)
  max_y = tf.to_float(max_y) / tf.to_float(image_height)
  max_x = tf.to_float(max_x) / tf.to_float(image_width)

  # Clip the bboxes to be sure the fall between [0, 1].
  min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
  min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
  return tf.stack([min_y, min_x, max_y, max_x])
//  [python]
def _apply_bbox_augmentation(
    image,
    bbox,
    augmentation_func,
    *args):
  """Applies augmentation_func to the subsection of image indicated by bbox.

  Args:
    image: 3D uint8 Tensor.
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    augmentation_func: Augmentation function that will be applied to the
      subsection of image.
    *args: Additional parameters that will be passed into augmentation_func
      when it is called.

  Returns:
    A modified version of image, where the bbox location in the image will
    have `ugmentation_func applied to it.
  """
  image_height = tf.to_float(tf.shape(image)[0])
  image_width = tf.to_float(tf.shape(image)[1])
  min_y = tf.to_int32(image_height * bbox[0])
  min_x = tf.to_int32(image_width * bbox[1])
  max_y = tf.to_int32(image_height * bbox[2])
  max_x = tf.to_int32(image_width * bbox[3])
  image_height = tf.to_int32(image_height)
  image_width = tf.to_int32(image_width)

  # Clip to be sure the max values do not fall out of range.
  max_y = tf.minimum(max_y, image_height - 1)
  max_x = tf.minimum(max_x, image_width - 1)

  # Get the sub-tensor that is the image within the bounding box region.
  bbox_content = image[min_y:max_y + 1, min_x:max_x + 1, :]

  # Apply the augmentation function to the bbox portion of the image.
  augmented_bbox_content = augmentation_func(bbox_content, *args)

  # Pad the augmented_bbox_content and the mask to match the shape of original
  # image.
  augmented_bbox_content = tf.pad(augmented_bbox_content,
                                  [[min_y, (image_height - 1) - max_y],
                                   [min_x, (image_width - 1) - max_x],
                                   [0, 0]])

  # Create a mask that will be used to zero out a part of the original image.
  mask_tensor = tf.zeros_like(bbox_content)

  mask_tensor = tf.pad(mask_tensor,
                       [[min_y, (image_height - 1) - max_y],
                        [min_x, (image_width - 1) - max_x],
                        [0, 0]],
                       constant_values=1)
  # Replace the old bbox content with the new augmented content.
  image = image * mask_tensor + augmented_bbox_content
  return image
//  [python]
def _apply_bbox_augmentation(
    image,
    bbox,
    augmentation_func,
    *args):
  """Applies augmentation_func to the subsection of image indicated by bbox.

  Args:
    image: 3D uint8 Tensor.
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    augmentation_func: Augmentation function that will be applied to the
      subsection of image.
    *args: Additional parameters that will be passed into augmentation_func
      when it is called.

  Returns:
    A modified version of image, where the bbox location in the image will
    have `ugmentation_func applied to it.
  """
  image_height = tf.cast(tf.shape(image)[0], tf.float32)
  image_width = tf.cast(tf.shape(image)[1], tf.float32)
  min_y = tf.to_int32(image_height * bbox[0])
  min_x = tf.to_int32(image_width * bbox[1])
  max_y = tf.to_int32(image_height * bbox[2])
  max_x = tf.to_int32(image_width * bbox[3])
  image_height = tf.to_int32(image_height)
  image_width = tf.to_int32(image_width)

  # Clip to be sure the max values do not fall out of range.
  max_y = tf.minimum(max_y, image_height - 1)
  max_x = tf.minimum(max_x, image_width - 1)

  # Get the sub-tensor that is the image within the bounding box region.
  bbox_content = image[min_y:max_y + 1, min_x:max_x + 1, :]

  # Apply the augmentation function to the bbox portion of the image.
  augmented_bbox_content = augmentation_func(bbox_content, *args)

  # Pad the augmented_bbox_content and the mask to match the shape of original
  # image.
  augmented_bbox_content = tf.pad(augmented_bbox_content,
                                  [[min_y, (image_height - 1) - max_y],
                                   [min_x, (image_width - 1) - max_x],
                                   [0, 0]])

  # Create a mask that will be used to zero out a part of the original image.
  mask_tensor = tf.zeros_like(bbox_content)

  mask_tensor = tf.pad(mask_tensor,
                       [[min_y, (image_height - 1) - max_y],
                        [min_x, (image_width - 1) - max_x],
                        [0, 0]],
                       constant_values=1)
  # Replace the old bbox content with the new augmented content.
  image = image * mask_tensor + augmented_bbox_content
  return image
//  [python]
def _shear_bbox(bbox,
                image_height,
                image_width,
                level,
                shear_horizontal):
  """Shifts the bbox according to how the image was sheared.

  Args:
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    image_height: Int, height of the image.
    image_width: Int, height of the image.
    level: Float. How much to shear the image.
    shear_horizontal: If true then shear in X dimension else shear in
      the Y dimension.

  Returns:
    A tensor of the same shape as bbox, but now with the shifted coordinates.
  """
  image_height, image_width = (
      tf.to_float(image_height), tf.to_float(image_width))

  # Change bbox coordinates to be pixels.
  min_y = tf.to_int32(image_height * bbox[0])
  min_x = tf.to_int32(image_width * bbox[1])
  max_y = tf.to_int32(image_height * bbox[2])
  max_x = tf.to_int32(image_width * bbox[3])
  coordinates = tf.stack(
      [[min_y, min_x], [min_y, max_x], [max_y, min_x], [max_y, max_x]])
  coordinates = tf.cast(coordinates, tf.float32)

  # Shear the coordinates according to the translation matrix.
  if shear_horizontal:
    translation_matrix = tf.stack(
        [[1, 0], [-level, 1]])
  else:
    translation_matrix = tf.stack(
        [[1, -level], [0, 1]])
  translation_matrix = tf.cast(translation_matrix, tf.float32)
  new_coords = tf.cast(
      tf.matmul(translation_matrix, tf.transpose(coordinates)), tf.int32)

  # Find min/max values and convert them back to floats.
  min_y = tf.to_float(tf.reduce_min(new_coords[0, :])) / image_height
  min_x = tf.to_float(tf.reduce_min(new_coords[1, :])) / image_width
  max_y = tf.to_float(tf.reduce_max(new_coords[0, :])) / image_height
  max_x = tf.to_float(tf.reduce_max(new_coords[1, :])) / image_width

  # Clip the bboxes to be sure the fall between [0, 1].
  min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
  min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
  return tf.stack([min_y, min_x, max_y, max_x])
//  [python]
def _shear_bbox(bbox,
                image_height,
                image_width,
                level,
                shear_horizontal):
  """Shifts the bbox according to how the image was sheared.

  Args:
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    image_height: Int, height of the image.
    image_width: Int, height of the image.
    level: Float. How much to shear the image.
    shear_horizontal: If true then shear in X dimension else shear in
      the Y dimension.

  Returns:
    A tensor of the same shape as bbox, but now with the shifted coordinates.
  """
  image_height, image_width = (
      tf.to_float(image_height), tf.to_float(image_width))

  # Change bbox coordinates to be pixels.
  min_y = tf.cast(image_height * bbox[0], tf.int32)
  min_x = tf.cast(image_width * bbox[1], tf.int32)
  max_y = tf.cast(image_height * bbox[2], tf.int32)
  max_x = tf.cast(image_width * bbox[3], tf.int32)
  coordinates = tf.stack(
      [[min_y, min_x], [min_y, max_x], [max_y, min_x], [max_y, max_x]])
  coordinates = tf.cast(coordinates, tf.float32)

  # Shear the coordinates according to the translation matrix.
  if shear_horizontal:
    translation_matrix = tf.stack(
        [[1, 0], [-level, 1]])
  else:
    translation_matrix = tf.stack(
        [[1, -level], [0, 1]])
  translation_matrix = tf.cast(translation_matrix, tf.float32)
  new_coords = tf.cast(
      tf.matmul(translation_matrix, tf.transpose(coordinates)), tf.int32)

  # Find min/max values and convert them back to floats.
  min_y = tf.to_float(tf.reduce_min(new_coords[0, :])) / image_height
  min_x = tf.to_float(tf.reduce_min(new_coords[1, :])) / image_width
  max_y = tf.to_float(tf.reduce_max(new_coords[0, :])) / image_height
  max_x = tf.to_float(tf.reduce_max(new_coords[1, :])) / image_width

  # Clip the bboxes to be sure the fall between [0, 1].
  min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
  min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
  return tf.stack([min_y, min_x, max_y, max_x])
//  [python]
def _rotate_bbox(bbox,
                 image_height,
                 image_width,
                 degrees):
  """Rotates the bbox coordinated by degrees.

  Args:
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    image_height: Int, height of the image.
    image_width: Int, height of the image.
    degrees: Float, a scalar angle in degrees to rotate all images by. If
      degrees is positive the image will be rotated clockwise otherwise it will
      be rotated counterclockwise.

  Returns:
    A tensor of the same shape as bbox, but now with the rotated coordinates.
  """
  image_height, image_width = (
      tf.to_float(image_height), tf.to_float(image_width))

  # Convert from degrees to radians.
  degrees_to_radians = math.pi / 180.0
  radians = degrees * degrees_to_radians

  # Translate the bbox to the center of the image and turn the normalized 0-1
  # coordinates to absolute pixel locations.
  # Y coordinates are made negative as the y axis of images goes down with
  # increasing pixel values, so we negate to make sure x axis and y axis points
  # are in the traditionally positive direction.
  min_y = -tf.to_int32(image_height * (bbox[0] - 0.5))
  min_x = tf.to_int32(image_width * (bbox[1] - 0.5))
  max_y = -tf.to_int32(image_height * (bbox[2] - 0.5))
  max_x = tf.to_int32(image_width * (bbox[3] - 0.5))
  coordinates = tf.stack(
      [[min_y, min_x], [min_y, max_x], [max_y, min_x], [max_y, max_x]])
  coordinates = tf.cast(coordinates, tf.float32)
  # Rotate the coordinates according to the rotation matrix clockwise if
  # radians is positive, else negative
  rotation_matrix = tf.stack(
      [[tf.cos(radians), tf.sin(radians)],
       [-tf.sin(radians), tf.cos(radians)]])
  new_coords = tf.cast(
      tf.matmul(rotation_matrix, tf.transpose(coordinates)), tf.int32)
  # Find min/max values and convert them back to normalized 0-1 floats.
  min_y = -(tf.to_float(tf.reduce_max(new_coords[0, :])) / image_height - 0.5)
  min_x = tf.to_float(tf.reduce_min(new_coords[1, :])) / image_width + 0.5
  max_y = -(tf.to_float(tf.reduce_min(new_coords[0, :])) / image_height - 0.5)
  max_x = tf.to_float(tf.reduce_max(new_coords[1, :])) / image_width + 0.5

  # Clip the bboxes to be sure the fall between [0, 1].
  min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
  min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
  return tf.stack([min_y, min_x, max_y, max_x])
//  [python]
def _rotate_bbox(bbox,
                 image_height,
                 image_width,
                 degrees):
  """Rotates the bbox coordinated by degrees.

  Args:
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    image_height: Int, height of the image.
    image_width: Int, height of the image.
    degrees: Float, a scalar angle in degrees to rotate all images by. If
      degrees is positive the image will be rotated clockwise otherwise it will
      be rotated counterclockwise.

  Returns:
    A tensor of the same shape as bbox, but now with the rotated coordinates.
  """
  image_height, image_width = (
      tf.to_float(image_height), tf.to_float(image_width))

  # Convert from degrees to radians.
  degrees_to_radians = math.pi / 180.0
  radians = degrees * degrees_to_radians

  # Translate the bbox to the center of the image and turn the normalized 0-1
  # coordinates to absolute pixel locations.
  # Y coordinates are made negative as the y axis of images goes down with
  # increasing pixel values, so we negate to make sure x axis and y axis points
  # are in the traditionally positive direction.
  min_y = -tf.cast(image_height * (bbox[0] - 0.5), tf.int32)
  min_x = tf.cast(image_width * (bbox[1] - 0.5), tf.int32)
  max_y = -tf.cast(image_height * (bbox[2] - 0.5), tf.int32)
  max_x = tf.cast(image_width * (bbox[3] - 0.5), tf.int32)
  coordinates = tf.stack(
      [[min_y, min_x], [min_y, max_x], [max_y, min_x], [max_y, max_x]])
  coordinates = tf.cast(coordinates, tf.float32)
  # Rotate the coordinates according to the rotation matrix clockwise if
  # radians is positive, else negative
  rotation_matrix = tf.stack(
      [[tf.cos(radians), tf.sin(radians)],
       [-tf.sin(radians), tf.cos(radians)]])
  new_coords = tf.cast(
      tf.matmul(rotation_matrix, tf.transpose(coordinates)), tf.int32)
  # Find min/max values and convert them back to normalized 0-1 floats.
  min_y = -(tf.to_float(tf.reduce_max(new_coords[0, :])) / image_height - 0.5)
  min_x = tf.to_float(tf.reduce_min(new_coords[1, :])) / image_width + 0.5
  max_y = -(tf.to_float(tf.reduce_min(new_coords[0, :])) / image_height - 0.5)
  max_x = tf.to_float(tf.reduce_max(new_coords[1, :])) / image_width + 0.5

  # Clip the bboxes to be sure the fall between [0, 1].
  min_y, min_x, max_y, max_x = _clip_bbox(min_y, min_x, max_y, max_x)
  min_y, min_x, max_y, max_x = _check_bbox_area(min_y, min_x, max_y, max_x)
  return tf.stack([min_y, min_x, max_y, max_x])
//  [python]
def loss_1(self, alpha):
        def loss_1st(y_true, y_pred):
            L = y_true
            Y = y_pred
            batch_size = tf.to_float(K.shape(L)[0])
            return alpha * 2 * tf.linalg.trace(tf.matmul(tf.matmul(Y, L, transpose_a=True), Y)) / batch_size
        return loss_1st
//  [python]
def loss_1(self, alpha):
        def loss_1st(y_true, y_pred):
            L = y_true
            Y = y_pred
            batch_size = tf.cast(K.shape(L)[0], tf.float32)
            return alpha * 2 * tf.linalg.trace(tf.matmul(tf.matmul(Y, L, transpose_a=True), Y)) / batch_size
        return loss_1st
//  [python]
def normalize(data):
  data['image'] = tf.to_float(data['image']) / 255.
  return data
//  [python]
def normalize(data):
  data['image'] = tf.cast(data['image'], tf.float32) / 255.
  return data
//  [python]
def __call__(self, features, is_training):
        seqs_id = features['seqs_i']
        seqs_inputs = self.item_embs(seqs_id)

        # Dropout
        seqs_units = tf.layers.dropout(seqs_inputs, rate=self.hidden_dropout_rate,
                                       training=tf.convert_to_tensor(is_training))
        seqs_masks = tf.expand_dims(tf.to_float(tf.not_equal(seqs_id, 0)), -1)

        # Recurrency
        with tf.variable_scope("S2PNM/Reccurency"):
            h, _ = self._cudnn_rnn(tf.transpose(seqs_units, [1, 0, 2]))
            h = tf.transpose(h, [1, 0, 2])  # mask the hidden states as dynamic_rnn did

        # Position coding and Mask
        seqs_units = self.pcoding(h)
        seqs_units = seqs_units * seqs_masks

        # Attention
        with tf.variable_scope("S2PNM/Attention"):
            seqs_units = self.attention(layernorm(seqs_units), seqs_units, is_training, causality=True)
        with tf.variable_scope("S2PNM/fforward"):
            g = self.fforward(layernorm(seqs_units), is_training)

        # Dictionary Learning
        with tf.variable_scope("S2PNM/Dictionary"):
            seqs_outs = layernorm(tf.concat([g, h, g - h, g * h], axis=-1))
            seqs_outs = tf.layers.dense(seqs_outs, 2 * self.num_units, activation=tf.nn.sigmoid)
            seqs_outs = tf.layers.dense(seqs_outs, self.num_units)
            seqs_outs += seqs_inputs

        if is_training:
            seqs_outs = tf.reshape(seqs_outs, [tf.shape(seqs_id)[0] * self.seqslen, self.num_units])
        else:
            # only using the latest representations for making predictions
            seqs_outs = tf.reshape(seqs_outs[:, -1], [tf.shape(seqs_id)[0], self.num_units])

        # compute logits
        logits = tf.matmul(seqs_outs, self.item_embs.lookup_table, transpose_b=True)
        logits = tf.nn.bias_add(logits, self.output_bias)
        return logits
//  [python]
def __call__(self, features, is_training):
        seqs_id = features['seqs_i']
        seqs_inputs = self.item_embs(seqs_id)

        # Dropout
        seqs_units = tf.layers.dropout(seqs_inputs, rate=self.hidden_dropout_rate,
                                       training=tf.convert_to_tensor(is_training))
        seqs_masks = tf.expand_dims(tf.cast(tf.not_equal(seqs_id, 0), tf.float32), -1)

        # Recurrency
        with tf.variable_scope("S2PNM/Reccurency"):
            h, _ = self._cudnn_rnn(tf.transpose(seqs_units, [1, 0, 2]))
            h = tf.transpose(h, [1, 0, 2])  # mask the hidden states as dynamic_rnn did

        # Position coding and Mask
        seqs_units = self.pcoding(h)
        seqs_units = seqs_units * seqs_masks

        # Attention
        with tf.variable_scope("S2PNM/Attention"):
            seqs_units = self.attention(layernorm(seqs_units), seqs_units, is_training, causality=True)
        with tf.variable_scope("S2PNM/fforward"):
            g = self.fforward(layernorm(seqs_units), is_training)

        # Dictionary Learning
        with tf.variable_scope("S2PNM/Dictionary"):
            seqs_outs = layernorm(tf.concat([g, h, g - h, g * h], axis=-1))
            seqs_outs = tf.layers.dense(seqs_outs, 2 * self.num_units, activation=tf.nn.sigmoid)
            seqs_outs = tf.layers.dense(seqs_outs, self.num_units)
            seqs_outs += seqs_inputs

        if is_training:
            seqs_outs = tf.reshape(seqs_outs, [tf.shape(seqs_id)[0] * self.seqslen, self.num_units])
        else:
            # only using the latest representations for making predictions
            seqs_outs = tf.reshape(seqs_outs[:, -1], [tf.shape(seqs_id)[0], self.num_units])

        # compute logits
        logits = tf.matmul(seqs_outs, self.item_embs.lookup_table, transpose_b=True)
        logits = tf.nn.bias_add(logits, self.output_bias)
        return logits
//  [python]
def __call__(self, queries, keys, users):
        num_units = queries.get_shape().as_list()[-1]

        with tf.variable_scope(self.scope):
            # N: batch_size, S: sequence_length, H: num_units, W: window_sizes
            users = tf.layers.dense(users, num_units, use_bias=False)  # N, S, H
            users = tf.expand_dims(users, axis=2)  # N, S, 1, H

            Q = queries * users  # N, S, 1, H
            K = keys * users  # N, S, W, H
            V = K  # N, S, W, H

            outputs = tf.matmul(Q, K, transpose_b=True)  # N, S, 1, W
            outputs = outputs / tf.sqrt(tf.to_float(num_units))
            outputs = tf.nn.softmax(outputs)  # N, S, 1, W

            outputs = tf.matmul(outputs, V)  # N, S, 1, H
            return tf.squeeze(outputs, axis=2)
//  [python]
def __call__(self, queries, keys, users):
        num_units = queries.get_shape().as_list()[-1]

        with tf.variable_scope(self.scope):
            # N: batch_size, S: sequence_length, H: num_units, W: window_sizes
            users = tf.layers.dense(users, num_units, use_bias=False)  # N, S, H
            users = tf.expand_dims(users, axis=2)  # N, S, 1, H

            Q = queries * users  # N, S, 1, H
            K = keys * users  # N, S, W, H
            V = K  # N, S, W, H

            outputs = tf.matmul(Q, K, transpose_b=True)  # N, S, 1, W
            outputs = outputs / tf.sqrt(tf.cast(num_units, tf.float32))
            outputs = tf.nn.softmax(outputs)  # N, S, 1, W

            outputs = tf.matmul(outputs, V)  # N, S, 1, H
            return tf.squeeze(outputs, axis=2)
//  [python]
def parse_example(serialized_example):
  """Parse example."""
  features = tf.parse_single_example(
      serialized_example,
      features={
          "question":
              tf.FixedLenFeature([], tf.string),
          "context":
              tf.FixedLenSequenceFeature(
                  dtype=tf.string, shape=[], allow_missing=True),
          "long_answer_indices":
              tf.FixedLenSequenceFeature(
                  dtype=tf.int64, shape=[], allow_missing=True)
      })
  features["question"] = features["question"]
  features["context"] = features["context"]
  features["long_answer_indices"] = tf.to_int32(features["long_answer_indices"])
  return features
//  [python]
def parse_example(serialized_example):
  """Parse example."""
  features = tf.parse_single_example(
      serialized_example,
      features={
          "question":
              tf.FixedLenFeature([], tf.string),
          "context":
              tf.FixedLenSequenceFeature(
                  dtype=tf.string, shape=[], allow_missing=True),
          "long_answer_indices":
              tf.FixedLenSequenceFeature(
                  dtype=tf.int64, shape=[], allow_missing=True)
      })
  features["question"] = features["question"]
  features["context"] = features["context"]
  features["long_answer_indices"] = tf.cast(features["long_answer_indices"], tf.int32)
  return features
//  [python]
def _cutout_inside_bbox(image,
                        bbox,
                        pad_fraction):
  """Generates cutout mask and the mean pixel value of the bbox.

  First a location is randomly chosen within the image as the center where the
  cutout mask will be applied. Note this can be towards the boundaries of the
  image, so the full cutout mask may not be applied.

  Args:
    image: 3D uint8 Tensor.
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    pad_fraction: Float that specifies how large the cutout mask should be in
      in reference to the size of the original bbox. If pad_fraction is 0.25,
      then the cutout mask will be of shape
      (0.25 * bbox height, 0.25 * bbox width).

  Returns:
    A tuple. Fist element is a tensor of the same shape as image where each
    element is either a 1 or 0 that is used to determine where the image
    will have cutout applied. The second element is the mean of the pixels
    in the image where the bbox is located.
  """
  image_height = tf.shape(image)[0]
  image_width = tf.shape(image)[1]
  # Transform from shape [1, 4] to [4].
  bbox = tf.squeeze(bbox)

  min_y = tf.to_int32(tf.to_float(image_height) * bbox[0])
  min_x = tf.to_int32(tf.to_float(image_width) * bbox[1])
  max_y = tf.to_int32(tf.to_float(image_height) * bbox[2])
  max_x = tf.to_int32(tf.to_float(image_width) * bbox[3])

  # Calculate the mean pixel values in the bounding box, which will be used
  # to fill the cutout region.
  mean = tf.reduce_mean(image[min_y:max_y + 1, min_x:max_x + 1],
                        reduction_indices=[0, 1])

  # Cutout mask will be size pad_size_heigh * 2 by pad_size_width * 2 if the
  # region lies entirely within the bbox.
  box_height = max_y - min_y + 1
  box_width = max_x - min_x + 1
  pad_size_height = tf.to_int32(pad_fraction * (box_height / 2))
  pad_size_width = tf.to_int32(pad_fraction * (box_width / 2))

  # Sample the center location in the image where the zero mask will be applied.
  cutout_center_height = tf.random_uniform(
      shape=[], minval=min_y, maxval=max_y+1,
      dtype=tf.int32)

  cutout_center_width = tf.random_uniform(
      shape=[], minval=min_x, maxval=max_x+1,
      dtype=tf.int32)

  lower_pad = tf.maximum(
      0, cutout_center_height - pad_size_height)
  upper_pad = tf.maximum(
      0, image_height - cutout_center_height - pad_size_height)
  left_pad = tf.maximum(
      0, cutout_center_width - pad_size_width)
  right_pad = tf.maximum(
      0, image_width - cutout_center_width - pad_size_width)

  cutout_shape = [image_height - (lower_pad + upper_pad),
                  image_width - (left_pad + right_pad)]
  padding_dims = [[lower_pad, upper_pad], [left_pad, right_pad]]

  mask = tf.pad(
      tf.zeros(cutout_shape, dtype=image.dtype),
      padding_dims, constant_values=1)

  mask = tf.expand_dims(mask, 2)
  mask = tf.tile(mask, [1, 1, 3])

  return mask, mean
//  [python]
def _cutout_inside_bbox(image,
                        bbox,
                        pad_fraction):
  """Generates cutout mask and the mean pixel value of the bbox.

  First a location is randomly chosen within the image as the center where the
  cutout mask will be applied. Note this can be towards the boundaries of the
  image, so the full cutout mask may not be applied.

  Args:
    image: 3D uint8 Tensor.
    bbox: 1D Tensor that has 4 elements (min_y, min_x, max_y, max_x)
      of type float that represents the normalized coordinates between 0 and 1.
    pad_fraction: Float that specifies how large the cutout mask should be in
      in reference to the size of the original bbox. If pad_fraction is 0.25,
      then the cutout mask will be of shape
      (0.25 * bbox height, 0.25 * bbox width).

  Returns:
    A tuple. Fist element is a tensor of the same shape as image where each
    element is either a 1 or 0 that is used to determine where the image
    will have cutout applied. The second element is the mean of the pixels
    in the image where the bbox is located.
  """
  image_height = tf.shape(image)[0]
  image_width = tf.shape(image)[1]
  # Transform from shape [1, 4] to [4].
  bbox = tf.squeeze(bbox)

  min_y = tf.cast(tf.to_float(image_height) * bbox[0], tf.int32)
  min_x = tf.cast(tf.to_float(image_width) * bbox[1], tf.int32)
  max_y = tf.cast(tf.to_float(image_height) * bbox[2], tf.int32)
  max_x = tf.cast(tf.to_float(image_width) * bbox[3], tf.int32)

  # Calculate the mean pixel values in the bounding box, which will be used
  # to fill the cutout region.
  mean = tf.reduce_mean(image[min_y:max_y + 1, min_x:max_x + 1],
                        reduction_indices=[0, 1])

  # Cutout mask will be size pad_size_heigh * 2 by pad_size_width * 2 if the
  # region lies entirely within the bbox.
  box_height = max_y - min_y + 1
  box_width = max_x - min_x + 1
  pad_size_height = tf.cast(pad_fraction * (box_height / 2), tf.int32)
  pad_size_width = tf.cast(pad_fraction * (box_width / 2), tf.int32)

  # Sample the center location in the image where the zero mask will be applied.
  cutout_center_height = tf.random_uniform(
      shape=[], minval=min_y, maxval=max_y+1,
      dtype=tf.int32)

  cutout_center_width = tf.random_uniform(
      shape=[], minval=min_x, maxval=max_x+1,
      dtype=tf.int32)

  lower_pad = tf.maximum(
      0, cutout_center_height - pad_size_height)
  upper_pad = tf.maximum(
      0, image_height - cutout_center_height - pad_size_height)
  left_pad = tf.maximum(
      0, cutout_center_width - pad_size_width)
  right_pad = tf.maximum(
      0, image_width - cutout_center_width - pad_size_width)

  cutout_shape = [image_height - (lower_pad + upper_pad),
                  image_width - (left_pad + right_pad)]
  padding_dims = [[lower_pad, upper_pad], [left_pad, right_pad]]

  mask = tf.pad(
      tf.zeros(cutout_shape, dtype=image.dtype),
      padding_dims, constant_values=1)

  mask = tf.expand_dims(mask, 2)
  mask = tf.tile(mask, [1, 1, 3])

  return mask, mean
//  [python]
def compute_pred_dict(candidates_dict, dev_features, raw_results):
  """Computes official answer key from raw logits."""
  sess = tf.Session()
  # raw_results_by_id = [
  # (int(res["unique_id"] + 1), res) for res in raw_results]
  raw_results = [res for res in raw_results]
  raw_results_ids = tf.to_int32(
      np.array([int(res["unique_id"] + 1) for res in raw_results
               ])).eval(session=sess)
  raw_results_by_id = list(zip(raw_results_ids, raw_results))

  # Cast example id to int32 for each example, similarly to the raw results.
  all_candidates = candidates_dict.items()
  example_ids = tf.to_int32(np.array([int(k) for k, _ in all_candidates
                                     ])).eval(session=sess)
  examples_by_id = list(zip(example_ids, all_candidates))

  # Cast unique_id also to int32 for features.
  feature_ids = []
  features = []
  for f in dev_features:
    feature_ids.append(f.features.feature["unique_ids"].int64_list.value[0] + 1)
    features.append(f.features.feature)
  feature_ids = tf.to_int32(np.array(feature_ids)).eval(session=sess)
  features_by_id = list(zip(feature_ids, features))

  # Join examplew with features and raw results.
  examples = []
  merged = sorted(
      examples_by_id + raw_results_by_id + features_by_id,
      key=lambda tup: tup[0])
  for idx, datum in merged:
    if isinstance(datum, tuple):
      examples.append(EvalExample(datum[0], datum[1]))
    elif "token_map" in datum:
      examples[-1].features[idx] = datum
    else:
      examples[-1].results[idx] = datum

  # Construct prediction objects.
  tf.logging.info("Computing predictions...")
  summary_dict = {}
  nq_pred_dict = {}
  for e in examples:
    summary = compute_predictions(e)
    summary_dict[e.example_id] = summary
    nq_pred_dict[e.example_id] = summary.predicted_label
    if len(nq_pred_dict) % 100 == 0:
      tf.logging.info("Examples processed: %d", len(nq_pred_dict))
  tf.logging.info("Done computing predictions.")

  return nq_pred_dict
//  [python]
def compute_pred_dict(candidates_dict, dev_features, raw_results):
  """Computes official answer key from raw logits."""
  sess = tf.Session()
  raw_results = [res for res in raw_results]
  raw_results_ids = tf.cast(
      np.array([int(res["unique_id"] + 1) for res in raw_results
               ]), tf.int32).eval(session=sess)
  raw_results_by_id = list(zip(raw_results_ids, raw_results))

  all_candidates = candidates_dict.items()
  example_ids = tf.cast(np.array([int(k) for k, _ in all_candidates
                                     ]), tf.int32).eval(session=sess)
  examples_by_id = list(zip(example_ids, all_candidates))

  feature_ids = []
  features = []
  for f in dev_features:
    feature_ids.append(f.features.feature["unique_ids"].int64_list.value[0] + 1)
    features.append(f.features.feature)
  feature_ids = tf.cast(np.array(feature_ids), tf.int32).eval(session=sess)
  features_by_id = list(zip(feature_ids, features))

  examples = []
  merged = sorted(
      examples_by_id + raw_results_by_id + features_by_id,
      key=lambda tup: tup[0])
  for idx, datum in merged:
    if isinstance(datum, tuple):
      examples.append(EvalExample(datum[0], datum[1]))
    elif "token_map" in datum:
      examples[-1].features[idx] = datum
    else:
      examples[-1].results[idx] = datum

  tf.logging.info("Computing predictions...")
  summary_dict = {}
  nq_pred_dict = {}
  for e in examples:
    summary = compute_predictions(e)
    summary_dict[e.example_id] = summary
    nq_pred_dict[e.example_id] = summary.predicted_label
    if len(nq_pred_dict) % 100 == 0:
      tf.logging.info("Examples processed: %d", len(nq_pred_dict))
  tf.logging.info("Done computing predictions.")

  return nq_pred_dict
//  [python]
def rotate_points_tensor(points, image, angle):

    s = tf.shape(image)
    image_center = tf.to_float(s[:2]) / 2.

    # center coordinates since rotation center is supposed to be in the image center
    points_centered = points - image_center

    rot_matrix = tf.dynamic_stitch([[0], [1], [2], [3]], [tf.cos(angle), -tf.sin(angle), tf.sin(angle), tf.cos(angle)])
    rot_matrix = tf.reshape(rot_matrix, shape=[2, 2])

    points_centered_rot = tf.matmul(rot_matrix, tf.transpose(points_centered))

    return tf.transpose(points_centered_rot) + image_center
//  [python]
def rotate_points_tensor(points, image, angle):

    s = tf.shape(image)
    image_center = tf.cast(s[:2], tf.float32) / 2.

    # center coordinates since rotation center is supposed to be in the image center
    points_centered = points - image_center

    rot_matrix = tf.dynamic_stitch([[0], [1], [2], [3]], [tf.cos(angle), -tf.sin(angle), tf.sin(angle), tf.cos(angle)])
    rot_matrix = tf.reshape(rot_matrix, shape=[2, 2])

    points_centered_rot = tf.matmul(rot_matrix, tf.transpose(points_centered))

    return tf.transpose(points_centered_rot) + image_center
//  [python]
def _image_from_feature(self, features):
        image = tf.image.decode_jpeg(features['image'], channels=3)
        image_height = tf.to_int32(features['height'])
        image_width = tf.to_int32(features['width'])
        image = tf.reshape(image, (image_height, image_width, 3))
        image = tf.to_float(image)
        return image, image_height, image_width
//  [python]
def _image_from_feature(self, features):
        image = tf.image.decode_jpeg(features['image'], channels=3)
        image_height = tf.cast(features['height'], tf.int32)
        image_width = tf.cast(features['width'], tf.int32)
        image = tf.reshape(image, (image_height, image_width, 3))
        image = tf.to_float(image)
        return image, image_height, image_width
//  [python]
def lms_to_heatmap(lms, h, w, n_landmarks, marked_index, sigma=5):
    xs, ys = tf.meshgrid(tf.range(0., tf.to_float(w)),
                         tf.range(0., tf.to_float(h)))
    gaussian = (1. / (sigma * np.sqrt(2. * np.pi)))
    marked_index = tf.to_int32(marked_index)

    def gaussian_fn(lms):
        y, x, idx = tf.unstack(lms)
        idx = tf.to_int32(idx)

        def run_true():
            return tf.exp(-0.5 * (tf.pow(ys - y, 2) + tf.pow(xs - x, 2)) *
                          tf.pow(1. / sigma, 2.)) * gaussian * 17.

        def run_false():
            return tf.zeros((h, w))

        return tf.cond(tf.reduce_any(tf.equal(marked_index, idx)), run_true, run_false)

    img_hm = tf.stack(tf.map_fn(gaussian_fn, tf.concat(
        [lms, tf.to_float(tf.range(0, n_landmarks))[..., None]], 1)))

    return img_hm
//  [python]
def lms_to_heatmap(lms, h, w, n_landmarks, marked_index, sigma=5):
    xs, ys = tf.meshgrid(tf.range(0., tf.to_float(w)),
                         tf.range(0., tf.to_float(h)))
    gaussian = (1. / (sigma * np.sqrt(2. * np.pi)))
    marked_index = tf.cast(marked_index, tf.int32)

    def gaussian_fn(lms):
        y, x, idx = tf.unstack(lms)
        idx = tf.cast(idx, tf.int32)

        def run_true():
            return tf.exp(-0.5 * (tf.pow(ys - y, 2) + tf.pow(xs - x, 2)) *
                          tf.pow(1. / sigma, 2.)) * gaussian * 17.

        def run_false():
            return tf.zeros((h, w))

        return tf.cond(tf.reduce_any(tf.equal(marked_index, idx)), run_true, run_false)

    img_hm = tf.stack(tf.map_fn(gaussian_fn, tf.concat(
        [lms, tf.cast(tf.range(0, n_landmarks), tf.float32)[..., None]], 1)))

    return img_hm
//  [python]
def _heatmap_from_feature(self, features):
        n_landmarks = tf.to_int32(features['n_landmarks'])
        gt_lms = tf.decode_raw(features['gt_pts'], tf.float32)
        mask_index = tf.decode_raw(features['mask_index'], tf.float32)
        gt_mask = tf.decode_raw(features['gt_mask'], tf.float32)
        image_height = tf.to_int32(features['height'])
        image_width = tf.to_int32(features['width'])

        gt_lms = tf.reshape(gt_lms, (n_landmarks, 2))
        gt_heatmap = lms_to_heatmap(
            gt_lms, image_height, image_width, n_landmarks, mask_index)
        gt_heatmap = tf.transpose(gt_heatmap, perm=[1,2,0])

        return gt_heatmap, gt_lms, n_landmarks, mask_index, gt_mask
//  [python]
def _heatmap_from_feature(self, features):
        n_landmarks = tf.cast(features['n_landmarks'], tf.int32)
        gt_lms = tf.decode_raw(features['gt_pts'], tf.float32)
        mask_index = tf.decode_raw(features['mask_index'], tf.float32)
        gt_mask = tf.decode_raw(features['gt_mask'], tf.float32)
        image_height = tf.cast(features['height'], tf.int32)
        image_width = tf.cast(features['width'], tf.int32)

        gt_lms = tf.reshape(gt_lms, (n_landmarks, 2))
        gt_heatmap = lms_to_heatmap(
            gt_lms, image_height, image_width, n_landmarks, mask_index)
        gt_heatmap = tf.transpose(gt_heatmap, perm=[1,2,0])

        return gt_heatmap, gt_lms, n_landmarks, mask_index, gt_mask
//  [python]
def _compute_new_dynamic_size(image, min_dimension, max_dimension):
    """Compute new dynamic shape for resize_to_range method."""
    image_shape = tf.shape(image)
    orig_height = tf.to_float(image_shape[0])
    orig_width = tf.to_float(image_shape[1])
    num_channels = image_shape[2]
    orig_min_dim = tf.minimum(orig_height, orig_width)
    # Calculates the larger of the possible sizes
    min_dimension = tf.constant(min_dimension, dtype=tf.float32)
    large_scale_factor = min_dimension / orig_min_dim
    # Scaling orig_(height|width) by large_scale_factor will make the smaller
    # dimension equal to min_dimension, save for floating point rounding errors.
    # For reasonably-sized images, taking the nearest integer will reliably
    # eliminate this error.
    large_height = tf.to_int32(tf.round(orig_height * large_scale_factor))
    large_width = tf.to_int32(tf.round(orig_width * large_scale_factor))
    large_size = tf.stack([large_height, large_width])
    if max_dimension:
        # Calculates the smaller of the possible sizes, use that if the larger
        # is too big.
        orig_max_dim = tf.maximum(orig_height, orig_width)
        max_dimension = tf.constant(max_dimension, dtype=tf.float32)
        small_scale_factor = max_dimension / orig_max_dim
        # Scaling orig_(height|width) by small_scale_factor will make the larger
        # dimension equal to max_dimension, save for floating point rounding
        # errors. For reasonably-sized images, taking the nearest integer will
        # reliably eliminate this error.
        small_height = tf.to_int32(tf.round(orig_height * small_scale_factor))
        small_width = tf.to_int32(tf.round(orig_width * small_scale_factor))
        small_size = tf.stack([small_height, small_width])
        new_size = tf.cond(
            tf.to_float(tf.reduce_max(large_size)) > max_dimension,
            lambda: small_size,
            lambda: large_size,
        )
    else:
        new_size = large_size
    return tf.stack(tf.unstack(new_size) + [num_channels])
//  [python]
def _compute_new_dynamic_size(image, min_dimension, max_dimension):
    """Compute new dynamic shape for resize_to_range method."""
    image_shape = tf.shape(image)
    orig_height = tf.cast(image_shape[0], tf.float32)
    orig_width = tf.cast(image_shape[1], tf.float32)
    num_channels = image_shape[2]
    orig_min_dim = tf.minimum(orig_height, orig_width)
    # Calculates the larger of the possible sizes
    min_dimension = tf.constant(min_dimension, dtype=tf.float32)
    large_scale_factor = min_dimension / orig_min_dim
    # Scaling orig_(height|width) by large_scale_factor will make the smaller
    # dimension equal to min_dimension, save for floating point rounding errors.
    # For reasonably-sized images, taking the nearest integer will reliably
    # eliminate this error.
    large_height = tf.to_int32(tf.round(orig_height * large_scale_factor))
    large_width = tf.to_int32(tf.round(orig_width * large_scale_factor))
    large_size = tf.stack([large_height, large_width])
    if max_dimension:
        # Calculates the smaller of the possible sizes, use that if the larger
        # is too big.
        orig_max_dim = tf.maximum(orig_height, orig_width)
        max_dimension = tf.constant(max_dimension, dtype=tf.float32)
        small_scale_factor = max_dimension / orig_max_dim
        # Scaling orig_(height|width) by small_scale_factor will make the larger
        # dimension equal to max_dimension, save for floating point rounding
        # errors. For reasonably-sized images, taking the nearest integer will
        # reliably eliminate this error.
        small_height = tf.to_int32(tf.round(orig_height * small_scale_factor))
        small_width = tf.to_int32(tf.round(orig_width * small_scale_factor))
        small_size = tf.stack([small_height, small_width])
        new_size = tf.cond(
            tf.to_float(tf.reduce_max(large_size)) > max_dimension,
            lambda: small_size,
            lambda: large_size,
        )
    else:
        new_size = large_size
    return tf.stack(tf.unstack(new_size) + [num_channels])
//  [python]
def calc_high_prob_overlaps(labels, probs, weights):
  """Calculates if the highest probability prediction is correct."""
  labels = tf.to_float(labels)
  probs *= weights
  labels *= weights
  probs = probs[:, :, :, 1]
  labels = labels[:, :, :, 0]

  overlap_probs = probs * labels
  non_overlap_probs = probs - overlap_probs

  return tf.greater(
      tf.reduce_max(overlap_probs, [1, 2]),
      tf.reduce_max(non_overlap_probs, [1, 2]))
//  [python]
def calc_high_prob_overlaps(labels, probs, weights):
  """Calculates if the highest probability prediction is correct."""
  labels = tf.cast(labels, tf.float32)
  probs *= weights
  labels *= weights
  probs = probs[:, :, :, 1]
  labels = labels[:, :, :, 0]

  overlap_probs = probs * labels
  non_overlap_probs = probs - overlap_probs

  return tf.greater(
      tf.reduce_max(overlap_probs, [1, 2]),
      tf.reduce_max(non_overlap_probs, [1, 2]))
//  [python]
def rotate_image_tensor(image, angle):
    s = tf.shape(image)
    image_center = tf.to_float(s[:2]) / 2.

    # Coordinates of new image
    xs, ys = tf.meshgrid(tf.range(0.,tf.to_float(s[1])), tf.range(0., tf.to_float(s[0])))
    coords_new = tf.reshape(tf.stack([ys,xs], 2), [-1, 2])

    # center coordinates since rotation center is supposed to be in the image center
    coords_new_centered = tf.to_float(coords_new) - image_center

    # Perform backward transformation of the image coordinates
    rot_mat_inv = tf.stack(
        [tf.cos(angle), tf.sin(angle), -tf.sin(angle), tf.cos(angle)])
    rot_mat_inv = tf.reshape(rot_mat_inv, shape=[2, 2])
    coord_old_centered = tf.matmul(
        rot_mat_inv, tf.transpose(coords_new_centered))
    coord_old = tf.to_int32(tf.round(
        tf.transpose(coord_old_centered) + image_center))


    # Find nearest neighbor in old image
    coord_old_y, coord_old_x = tf.unstack(coord_old, axis=1)


    # Clip values to stay inside image coordinates
    outside_y = tf.logical_or(tf.greater(
        coord_old_y, s[0]-1), tf.less(coord_old_y, 0))
    outside_x = tf.logical_or(tf.greater(
        coord_old_x, s[1]-1), tf.less(coord_old_x, 0))
    outside_ind = tf.logical_or(outside_y, outside_x)


    inside_mask = tf.logical_not(outside_ind)
    inside_mask = tf.tile(tf.reshape(inside_mask, s[:2])[...,None], tf.stack([1,1,s[2]]))

    coord_old_y = tf.clip_by_value(coord_old_y, 0, s[0]-1)
    coord_old_x = tf.clip_by_value(coord_old_x, 0, s[1]-1)
    coord_flat = coord_old_y * s[1] + coord_old_x

    im_flat = tf.reshape(image, tf.stack([-1, s[2]]))
    rot_image = tf.gather(im_flat, coord_flat)
    rot_image = tf.reshape(rot_image, s)


    return tf.where(inside_mask, rot_image, tf.zeros_like(rot_image))