jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def tf_batch_map_coordinates(inputs, coords):
    """Batch version of tf_map_coordinates

    Only supports 2D feature maps

    Parameters
    ----------
    input : tf.Tensor. shape = (b*c, h, w)
    coords : tf.Tensor. shape = (b*c, h, w, n, 2)

    Returns
    -------
    tf.Tensor. shape = (b*c, h, w, n)
    """

    input_shape = inputs.get_shape()
    coords_shape = coords.get_shape()
    batch_channel = tf.shape(inputs)[0]
    input_h = int(input_shape[1])
    input_w = int(input_shape[2])
    kernel_n = int(coords_shape[3])
    n_coords = input_h * input_w * kernel_n

    coords_lt = t.floor(coords)
    coords_rb = t.ceil(coords)
    coords_lb = tf.stack([coords_lt[:, :, :, :, 0], coords_rb[:, :, :, :, 1]], axis=-1)
    coords_rt = tf.stack([coords_rb[:, :, :, :, 0], coords_lt[:, :, :, :, 1]], axis=-1)

    idx = tf_repeat(tf.range(batch_channel), n_coords)

    vals_lt = _get_vals_by_coords(inputs, coords_lt, idx, (batch_channel, input_h, input_w, kernel_n))
    vals_rb = _get_vals_by_coords(inputs, coords_rb, idx, (batch_channel, input_h, input_w, kernel_n))
    vals_lb = _get_vals_by_coords(inputs, coords_lb, idx, (batch_channel, input_h, input_w, kernel_n))
    vals_rt = _get_vals_by_coords(inputs, coords_rt, idx, (batch_channel, input_h, input_w, kernel_n))

    coords_offset_lt = coords - t.cast(coords_lt, 'float32')

    vals_t = vals_lt + (vals_rt - vals_lt) * coords_offset_lt[:, :, :, :, 0]
    vals_b = vals_lb + (vals_rb - vals_lb) * coords_offset_lt[:, :, :, :, 0]
    mapped_vals = vals_t + (vals_b - vals_t) * coords_offset_lt[:, :, :, :, 1]

    return mapped_vals
//  [python]
def tf_batch_map_coordinates(inputs, coords):
    """Batch version of tf_map_coordinates

    Only supports 2D feature maps

    Parameters
    ----------
    input : tf.Tensor. shape = (b*c, h, w)
    coords : tf.Tensor. shape = (b*c, h, w, n, 2)

    Returns
    -------
    tf.Tensor. shape = (b*c, h, w, n)
    """

    input_shape = inputs.get_shape()
    coords_shape = coords.get_shape()
    batch_channel = tf.shape(inputs)[0]
    input_h = int(input_shape[1])
    input_w = int(input_shape[2])
    kernel_n = int(coords_shape[3])
    n_coords = input_h * input_w * kernel_n

    coords_lt = tf.cast(tf.floor(coords), 'int32')
    coords_rb = tf.cast(tf.ceil(coords), 'int32')
    coords_lb = tf.stack([coords_lt[:, :, :, :, 0], coords_rb[:, :, :, :, 1]], axis=-1)
    coords_rt = tf.stack([coords_rb[:, :, :, :, 0], coords_lt[:, :, :, :, 1]], axis=-1)

    idx = tf_repeat(tf.range(batch_channel), n_coords)

    vals_lt = _get_vals_by_coords(inputs, coords_lt, idx, (batch_channel, input_h, input_w, kernel_n))
    vals_rb = _get_vals_by_coords(inputs, coords_rb, idx, (batch_channel, input_h, input_w, kernel_n))
    vals_lb = _get_vals_by_coords(inputs, coords_lb, idx, (batch_channel, input_h, input_w, kernel_n))
    vals_rt = _get_vals_by_coords(inputs, coords_rt, idx, (batch_channel, input_h, input_w, kernel_n))

    coords_offset_lt = coords - tf.cast(coords_lt, 'float32')

    vals_t = vals_lt + (vals_rt - vals_lt) * coords_offset_lt[:, :, :, :, 0]
    vals_b = vals_lb + (vals_rb - vals_lb) * coords_offset_lt[:, :, :, :, 0]
    mapped_vals = vals_t + (vals_b - vals_t) * coords_offset_lt[:, :, :, :, 1]

    return mapped_vals
//  [python]
def random_rotation(images,
                    flow,
                    mask,
                    probability,
                    max_rotation,
                    not_empty_crop = True):
  """Performs a random rotation with the specified maximum rotation."""
  perform_rotation = tf.random.uniform([]) < probability
  def true_fn(images, flow, mask):

    angle_radian = tf.random.uniform(
        [], minval=-max_rotation, maxval=max_rotation,
        dtype=tf.float32) * math.pi / 180.0
    images = rotate(images, angle_radian, is_flow=False, mask=None)

    if not_empty_crop:
      orig_height = tf.shape(images)[-3]
      orig_width = tf.shape(images)[-2]
      # introduce abbreviations for shorter notation
      cos = tf.math.cos(angle_radian % math.pi)
      sin = tf.math.sin(angle_radian % math.pi)
      h = tf.to_float(orig_height)
      w = tf.to_float(orig_width)

      # compute required scale factor
      scale = tf.cond(tf.math.less(angle_radian % math.pi, math.pi/2.0),
                      lambda: tf.math.maximum((w/h)*sin+cos, (h/w)*sin+cos),
                      lambda: tf.math.maximum((w/h)*sin-cos, (h/w)*sin-cos))
      new_height = tf.math.floor(h / scale)
      new_width = tf.math.floor(w / scale)

      # crop image again to original size
      offset_height = tf.to_int32((h - new_height) / 2)
      offset_width = tf.to_int32((w - new_width) / 2)
      images = tf.image.crop_to_bounding_box(
          images,
          offset_height=offset_height,
          offset_width=offset_width,
          target_height=tf.to_int32(new_height),
          target_width=tf.to_int32(new_width))

    if flow is not None:
      flow, mask = rotate(flow, angle_radian, is_flow=True, mask=mask)

      if not_empty_crop:
        # crop flow and mask again to original size
        flow = tf.image.crop_to_bounding_box(
            flow,
            offset_height=offset_height,
            offset_width=offset_width,
            target_height=tf.to_int32(new_height),
            target_width=tf.to_int32(new_width))
        mask = tf.image.crop_to_bounding_box(
            mask,
            offset_height=offset_height,
            offset_width=offset_width,
            target_height=tf.to_int32(new_height),
            target_width=tf.to_int32(new_width))
    return images, flow, mask
  def false_fn(images, flow, mask):
    return images, flow, mask

  return tf.cond(perform_rotation, lambda: true_fn(images, flow, mask),
                 lambda: false_fn(images, flow, mask))
//  [python]
def random_rotation(images,
                    flow,
                    mask,
                    probability,
                    max_rotation,
                    not_empty_crop = True):
  """Performs a random rotation with the specified maximum rotation."""
  perform_rotation = tf.random.uniform([]) < probability
  def true_fn(images, flow, mask):

    angle_radian = tf.random.uniform(
        [], minval=-max_rotation, maxval=max_rotation,
        dtype=tf.float32) * math.pi / 180.0
    images = rotate(images, angle_radian, is_flow=False, mask=None)

    if not_empty_crop:
      orig_height = tf.shape(images)[-3]
      orig_width = tf.shape(images)[-2]
      # introduce abbreviations for shorter notation
      cos = tf.math.cos(angle_radian % math.pi)
      sin = tf.math.sin(angle_radian % math.pi)
      h = tf.cast(orig_height, tf.float32)
      w = tf.cast(orig_width, tf.float32)

      # compute required scale factor
      scale = tf.cond(tf.math.less(angle_radian % math.pi, math.pi/2.0),
                      lambda: tf.math.maximum((w/h)*sin+cos, (h/w)*sin+cos),
                      lambda: tf.math.maximum((w/h)*sin-cos, (h/w)*sin-cos))
      new_height = tf.math.floor(h / scale)
      new_width = tf.math.floor(w / scale)

      # crop image again to original size
      offset_height = tf.cast((h - new_height) / 2, tf.int32)
      offset_width = tf.cast((w - new_width) / 2, tf.int32)
      images = tf.image.crop_to_bounding_box(
          images,
          offset_height=offset_height,
          offset_width=offset_width,
          target_height=tf.cast(new_height, tf.int32),
          target_width=tf.cast(new_width, tf.int32))

    if flow is not None:
      flow, mask = rotate(flow, angle_radian, is_flow=True, mask=mask)

      if not_empty_crop:
        # crop flow and mask again to original size
        flow = tf.image.crop_to_bounding_box(
            flow,
            offset_height=offset_height,
            offset_width=offset_width,
            target_height=tf.cast(new_height, tf.int32),
            target_width=tf.cast(new_width, tf.int32))
        mask = tf.image.crop_to_bounding_box(
            mask,
            offset_height=offset_height,
            offset_width=offset_width,
            target_height=tf.cast(new_height, tf.int32),
            target_width=tf.cast(new_width, tf.int32))
    return images, flow, mask
  def false_fn(images, flow, mask):
    return images, flow, mask

  return tf.cond(perform_rotation, lambda: true_fn(images, flow, mask),
                 lambda: false_fn(images, flow, mask))
//  [python]
def random_rotation_second(images,
                           flow,
                           mask,
                           probability,
                           max_rotation,
                           not_empty_crop=True):
  """Performs a random rotation on only the second image."""
  perform_rotation = tf.random.uniform([]) < probability
  def true_fn(images, flow, mask):
    angle_radian = tf.random.uniform(
        [], minval=-max_rotation, maxval=max_rotation,
        dtype=tf.float32) * math.pi / 180.0

    image_1, image_2 = tf.unstack(images)
    image_2 = rotate(image_2, angle_radian, is_flow=False, mask=None)
    images = tf.stack([image_1, image_2])

    if not_empty_crop:
      orig_height = tf.shape(images)[-3]
      orig_width = tf.shape(images)[-2]
      # introduce abbreviations for shorter notation
      cos = tf.math.cos(angle_radian % math.pi)
      sin = tf.math.sin(angle_radian % math.pi)
      h = tf.to_float(orig_height)
      w = tf.to_float(orig_width)

      # compute required scale factor
      scale = tf.cond(tf.math.less(angle_radian % math.pi, math.pi/2.0),
                      lambda: tf.math.maximum((w/h)*sin+cos, (h/w)*sin+cos),
                      lambda: tf.math.maximum((w/h)*sin-cos, (h/w)*sin-cos))
      new_height = tf.math.floor(h / scale)
      new_width = tf.math.floor(w / scale)

      # crop image again to original size
      offset_height = tf.to_int32((h-new_height)/2)
      offset_width = tf.to_int32((w-new_width)/2)
      images = tf.image.crop_to_bounding_box(
          images,
          offset_height=offset_height,
          offset_width=offset_width,
          target_height=tf.to_int32(new_height),
          target_width=tf.to_int32(new_width))

    if flow is not None:
      # get current locations (with the origin in the image center)
      positions = _positions_center_origin(orig_height, orig_width)

      # compute augmented flow (multiply by mask to zero invalid flow locations)
      cos = tf.math.cos(angle_radian)
      sin = tf.math.sin(angle_radian)
      rotation_matrix = tf.reshape([cos, sin, -sin, cos], [2, 2])
      flow = (tf.linalg.matmul(
          (positions + flow), rotation_matrix) - positions) * mask

      if not_empty_crop:
        # crop flow and mask again to original size
        flow = tf.image.crop_to_bounding_box(
            flow,
            offset_height=offset_height,
            offset_width=offset_width,
            target_height=tf.to_int32(new_height),
            target_width=tf.to_int32(new_width))
        mask = tf.image.crop_to_bounding_box(
            mask,
            offset_height=offset_height,
            offset_width=offset_width,
            target_height=tf.to_int32(new_height),
            target_width=tf.to_int32(new_width))
    return images, flow, mask
  def false_fn(images, flow, mask):
    return images, flow, mask

  return tf.cond(perform_rotation, lambda: true_fn(images, flow, mask),
                 lambda: false_fn(images, flow, mask))
//  [python]
def random_rotation_second(images,
                           flow,
                           mask,
                           probability,
                           max_rotation,
                           not_empty_crop=True):
  """Performs a random rotation on only the second image."""
  perform_rotation = tf.random.uniform([]) < probability
  def true_fn(images, flow, mask):
    angle_radian = tf.random.uniform(
        [], minval=-max_rotation, maxval=max_rotation,
        dtype=tf.float32) * math.pi / 180.0

    image_1, image_2 = tf.unstack(images)
    image_2 = rotate(image_2, angle_radian, is_flow=False, mask=None)
    images = tf.stack([image_1, image_2])

    if not_empty_crop:
      orig_height = tf.shape(images)[-3]
      orig_width = tf.shape(images)[-2]
      # introduce abbreviations for shorter notation
      cos = tf.math.cos(angle_radian % math.pi)
      sin = tf.math.sin(angle_radian % math.pi)
      h = tf.cast(orig_height, tf.float32)
      w = tf.cast(orig_width, tf.float32)

      # compute required scale factor
      scale = tf.cond(tf.math.less(angle_radian % math.pi, math.pi/2.0),
                      lambda: tf.math.maximum((w/h)*sin+cos, (h/w)*sin+cos),
                      lambda: tf.math.maximum((w/h)*sin-cos, (h/w)*sin-cos))
      new_height = tf.math.floor(h / scale)
      new_width = tf.math.floor(w / scale)

      # crop image again to original size
      offset_height = tf.cast((h-new_height)/2, tf.int32)
      offset_width = tf.cast((w-new_width)/2, tf.int32)
      images = tf.image.crop_to_bounding_box(
          images,
          offset_height=offset_height,
          offset_width=offset_width,
          target_height=tf.cast(new_height, tf.int32),
          target_width=tf.cast(new_width, tf.int32))

    if flow is not None:
      # get current locations (with the origin in the image center)
      positions = _positions_center_origin(orig_height, orig_width)

      # compute augmented flow (multiply by mask to zero invalid flow locations)
      cos = tf.math.cos(angle_radian)
      sin = tf.math.sin(angle_radian)
      rotation_matrix = tf.reshape([cos, sin, -sin, cos], [2, 2])
      flow = (tf.linalg.matmul(
          (positions + flow), rotation_matrix) - positions) * mask

      if not_empty_crop:
        # crop flow and mask again to original size
        flow = tf.image.crop_to_bounding_box(
            flow,
            offset_height=offset_height,
            offset_width=offset_width,
            target_height=tf.cast(new_height, tf.int32),
            target_width=tf.cast(new_width, tf.int32))
        mask = tf.image.crop_to_bounding_box(
            mask,
            offset_height=offset_height,
            offset_width=offset_width,
            target_height=tf.cast(new_height, tf.int32),
            target_width=tf.cast(new_width, tf.int32))
    return images, flow, mask
  def false_fn(images, flow, mask):
    return images, flow, mask

  return tf.cond(perform_rotation, lambda: true_fn(images, flow, mask),
                 lambda: false_fn(images, flow, mask))
//  [python]
def target_mask_op(data, pad_val=0):  # HangSheng: return tensor for mask,if input is tf.string
    data_shape_size = data.get_shape().ndims
    if data_shape_size == 3:
        return tf.to_int32(tf.reduce_any(tf.not_equal(data, pad_val), axis=2))
    elif data_shape_size == 2:
        return tf.to_int32(tf.not_equal(data, pad_val))
    elif data_shape_size == 1:
        raise ValueError("target_mask_op: data has wrong shape!")
    else:
        raise ValueError("target_mask_op: handling data_shape_size %s hasn't been implemented!" % (data_shape_size))
//  [python]
def target_mask_op(data, pad_val=0):  # HangSheng: return tensor for mask,if input is tf.string
    data_shape_size = data.get_shape().ndims
    if data_shape_size == 3:
        return tf.cast(tf.reduce_any(tf.not_equal(data, pad_val), axis=2), dtype=tf.int32)
    elif data_shape_size == 2:
        return tf.cast(tf.not_equal(data, pad_val), dtype=tf.int32)
    elif data_shape_size == 1:
        raise ValueError("target_mask_op: data has wrong shape!")
    else:
        raise ValueError("target_mask_op: handling data_shape_size %s hasn't been implemented!" % (data_shape_size))
//  [python]
def call(self, inputs: t_input):
        # tf.print(f"inputs at ComplexFlatten are {inputs.dtype}")
        real_flat = super(ComplexFlatten, self).call(tf.math.real(inputs))
        imag_flat = super(ComplexFlatten, self).call(tf.math.imag(inputs))
        return tf.complex(real_flat, imag_flat).dtype(inputs.dtype)
//  [python]
def call(self, inputs: t_input):
        # tf.print(f"inputs at ComplexFlatten are {inputs.dtype}")
        real_flat = super(ComplexFlatten, self).call(tf.math.real(inputs))
        imag_flat = super(ComplexFlatten, self).call(tf.math.imag(inputs))
        return tf.cast(tf.complex(real_flat, imag_flat), inputs.dtype)
//  [python]
def parser(self, line):
        features = tf.parse_single_example(line, features=self.proto)
        features['token_ids'] = t.cast(features['token_ids'], t.int32)
        features['mask'] = t.cast(features['mask'], t.int32)
        features['segment_ids'] = t.cast(features['segment_ids'], t.int32)
        features['label_ids'] = t.cast(features['label_ids'], t.int32)
        features['seq_len'] = tf.squeeze(t.cast(features['seq_len'], t.int32))
        # adjust parser when word enhance method is used
        if self.surfix == SoftWord:
            features['softword_ids'] = t.cast(features['softword_ids'], t.int32)
        elif self.surfix == ExSoftWord:
            # cast to float and reshape to original 2 dimension
            features['ex_softword_ids'] = t.cast(features['ex_softword_ids'], t.float32)
        elif self.surfix == SoftLexicon:
            features['softlexicon_ids'] = t.cast(features['softlexicon_ids'], t.int32)
            features['softlexicon_weights'] = t.cast(features['softlexicon_weights'], t.float32)
        return features
//  [python]
def parser(self, line):
        features = tf.parse_single_example(line, features=self.proto)
        features['token_ids'] = tf.cast(features['token_ids'], tf.int32)
        features['mask'] = tf.cast(features['mask'], tf.int32)
        features['segment_ids'] = tf.cast(features['segment_ids'], tf.int32)
        features['label_ids'] = tf.cast(features['label_ids'], tf.int32)
        features['seq_len'] = tf.squeeze(tf.cast(features['seq_len'], tf.int32))
        # adjust parser when word enhance method is used
        if self.surfix == SoftWord:
            features['softword_ids'] = tf.cast(features['softword_ids'], tf.int32)
        elif self.surfix == ExSoftWord:
            # cast to float and reshape to original 2 dimension
            features['ex_softword_ids'] = tf.cast(features['ex_softword_ids'], tf.float32)
        elif self.surfix == SoftLexicon:
            features['softlexicon_ids'] = tf.cast(features['softlexicon_ids'], tf.int32)
            features['softlexicon_weights'] = tf.cast(features['softlexicon_weights'], tf.float32)
        return features
//  [python]
def call(self, inputs, training=None):
        """
        :param inputs: Input tensor (of any rank).
        :param training: Python boolean indicating whether the layer should behave in training mode (adding dropout)
            or in inference mode (doing nothing).
        """
        if training is None:
            training = K.learning_phase()
            tf.print(f"Training was None and now is {training}")
            # This is used for my own debugging, I don't know WHEN this happens,
            # I trust K.learning_phase() returns a correct boolean.

        # def dropped_inputs():
        #     # import pdb; pdb.set_trace()
        #     drop_filter = tf.nn.dropout(tf.ones(tf.shape(inputs)), rate=self.rate,
        #                                 noise_shape=self._get_noise_shape(inputs), seed=self.seed)
        #     y_out = tf.multiply(t.cast(drop_filter, dtype=inputs.dtype), inputs)
        #     y_out = t.cast(y_out, dtype=inputs.dtype)
        #     return y_out
        # output = control_flow_util.smart_cond(training, dropped_inputs, lambda: tf.identity(inputs))
        # return output
        if not training:
            return inputs
        drop_filter = tf.nn.dropout(tf.ones(tf.shape(inputs)), rate=self.rate,
                                    noise_shape=self.noise_shape, seed=self.seed)
        y_out = tf.multiply(t.cast(drop_filter, dtype=inputs.dtype), inputs)
        y_out = t.cast(y_out, dtype=inputs.dtype)
        return y_out
//  [python]
def call(self, inputs, training=None):
        """
        :param inputs: Input tensor (of any rank).
        :param training: Python boolean indicating whether the layer should behave in training mode (adding dropout)
            or in inference mode (doing nothing).
        """
        if training is None:
            training = K.learning_phase()
            tf.print(f"Training was None and now is {training}")
            # This is used for my own debugging, I don't know WHEN this happens,
            # I trust K.learning_phase() returns a correct boolean.

        # def dropped_inputs():
        #     # import pdb; pdb.set_trace()
        #     drop_filter = tf.nn.dropout(tf.ones(tf.shape(inputs)), rate=self.rate,
        #                                 noise_shape=self._get_noise_shape(inputs), seed=self.seed)
        #     y_out = tf.multiply(tf.cast(drop_filter, dtype=inputs.dtype), inputs)
        #     y_out = tf.cast(y_out, dtype=inputs.dtype)
        #     return y_out
        # output = control_flow_util.smart_cond(training, dropped_inputs, lambda: tf.identity(inputs))
        # return output
        if not training:
            return inputs
        drop_filter = tf.nn.dropout(tf.ones(tf.shape(inputs)), rate=self.rate,
                                    noise_shape=self.noise_shape, seed=self.seed)
        y_out = tf.multiply(tf.cast(drop_filter, dtype=inputs.dtype), inputs)
        y_out = tf.cast(y_out, dtype=inputs.dtype)
        return y_out
//  [python]
def preprocess_images(self, images):
    """Applies preprocessing operations for Inception images.

    Because this will run in model_fn, on the accelerator, we use operations
    that efficiently execute there.

    Args:
      images: An Tensor of shape [batch_size height, width, channel] with uint8
        values.

    Returns:
      A tensor of images of shape [batch_size height, width, channel]
      containing floating point values, with all points rescaled between
      -1 and 1 and possibly resized.
    """
    images = tf.to_float(images)
    images = tf.subtract(images, 128.0)
    images = tf.math.divide(images, 128.0)
    return images
//  [python]
def preprocess_images(self, images):
    """Applies preprocessing operations for Inception images.

    Because this will run in model_fn, on the accelerator, we use operations
    that efficiently execute there.

    Args:
      images: An Tensor of shape [batch_size height, width, channel] with uint8
        values.

    Returns:
      A tensor of images of shape [batch_size height, width, channel]
      containing floating point values, with all points rescaled between
      -1 and 1 and possibly resized.
    """
    images = tf.cast(images, dtype=tf.float32)
    images = tf.subtract(images, 128.0)
    images = tf.math.divide(images, 128.0)
    return images
//  [python]
def preprocess_images(self, images):
    """Preprocess images for placeholder model."""
    # Note these calculations aren't necessary, but they are included here to
    # mimic the data processing pipeline used by inception. We may consider
    # removing them in a future CL, or making them optional, to reduce CPU cost
    # of this model.
    images = tf.cast(images, dtype=tf.float32)
    images = tf.subtract(images, 128.0)
    images = tf.math.divide(images, 128.0)
    return images
//  [python]
def _normalize(self, inputs, var, mean):
        """
        :inputs: Tensor
        :param var: Tensor of shape [..., 2, 2], if inputs dtype is real, var[slice] = [[var_slice, 0], [0, 0]]
        :param mean: Tensor with the mean in the corresponding dtype (same shape as inputs)
        """
        complex_zero_mean = inputs - mean
        # Inv and sqrtm is done over 2 inner most dimension [..., M, M] so it should be [..., 2, 2] for us.
        inv_sqrt_var = tf.linalg.sqrtm(tf.linalg.inv(var + self.epsilon_matrix))  # var^(-1/2)  # TODO: Check this exists always?
        # Separate real and imag so I go from shape [...] to [..., 2]
        zero_mean = tf.stack((tf.math.real(complex_zero_mean), tf.math.imag(complex_zero_mean)), axis=-1)
        # I expand dims to make the mult of matrix [..., 2, 2] and [..., 2, 1] resulting in [..., 2, 1]
        inputs_hat = tf.matmul(inv_sqrt_var, tf.expand_dims(zero_mean, axis=-1))

        # Then I squeeze to remove the last shape so I go from [..., 2, 1] to [..., 2].
        # Use reshape and not squeeze in case I have 1 channel for example.
        squeeze_inputs_hat = tf.reshape(inputs_hat, shape=tf.shape(inputs_hat)[:-1])
        # Get complex data
        complex_inputs_hat = tf.complex(squeeze_inputs_hat[..., 0], squeeze_inputs_hat[..., 1])
        # import pdb; pdb.set_trace()
        return complex_inputs_hat
//  [python]
def _normalize(self, inputs, var, mean):
        """
        :inputs: Tensor
        :param var: Tensor of shape [..., 2, 2], if inputs dtype is real, var[slice] = [[var_slice, 0], [0, 0]]
        :param mean: Tensor with the mean in the corresponding dtype (same shape as inputs)
        """
        complex_zero_mean = inputs - mean
        # Inv and sqrtm is done over 2 inner most dimension [..., M, M] so it should be [..., 2, 2] for us.
        inv_sqrt_var = tf.linalg.sqrtm(tf.linalg.inv(var + self.epsilon_matrix))  # var^(-1/2)  # TODO: Check this exists always?
        # Separate real and imag so I go from shape [...] to [..., 2]
        zero_mean = tf.stack((tf.math.real(complex_zero_mean), tf.math.imag(complex_zero_mean)), axis=-1)
        # I expand dims to make the mult of matrix [..., 2, 2] and [..., 2, 1] resulting in [..., 2, 1]
        inputs_hat = tf.matmul(inv_sqrt_var, tf.expand_dims(zero_mean, axis=-1))

        # Then I squeeze to remove the last shape so I go from [..., 2, 1] to [..., 2].
        # Use reshape and not squeeze in case I have 1 channel for example.
        squeeze_inputs_hat = tf.reshape(inputs_hat, shape=tf.shape(inputs_hat)[:-1])
        # Get complex data
        complex_inputs_hat = tf.cast(tf.complex(squeeze_inputs_hat[..., 0], squeeze_inputs_hat[..., 1]),
                                     dtype=self.my_dtype)
        # import pdb; pdb.set_trace()
        return complex_inputs_hat
//  [python]
def _convert_string_vector_to_action_input(example):
  return (
      ts.TimeStep(
          step_type=tf.to_int32(
              tf.strings.to_number(example[:, 0], tf.float32)
          ),
          reward=tf.strings.to_number(example[:, 1], tf.float32),
          discount=tf.strings.to_number(example[:, 2], tf.float32),
          observation=tf.strings.to_number(example[:, 3:7], tf.float32),
      ),
      (),
  )
//  [python]
def _convert_string_vector_to_action_input(example):
  return (
      ts.TimeStep(
          step_type=tf.cast(
              tf.strings.to_number(example[:, 0], tf.float32), tf.int32
          ),
          reward=tf.strings.to_number(example[:, 1], tf.float32),
          discount=tf.strings.to_number(example[:, 2], tf.float32),
          observation=tf.strings.to_number(example[:, 3:7], tf.float32),
      ),
      (),
  )
//  [python]
def create_input_data(source_data_file, target_data_file,
                      source_vocab_file, target_vocab_file,
                      batch_size, unk_id, sos, eos,
                      source_max_length, target_max_length):
  source_dataset = tf.data.TextLineDataset(tf.gfile.Glob(source_data_file))
  target_dataset = tf.data.TextLineDataset(tf.gfile.Glob(target_data_file))
  source_vocab = lookup_ops.index_table_from_file(
    source_vocab_file, default_value=unk_id)
  target_vocab = lookup_ops.index_table_from_file(
    target_vocab_file, default_value=unk_id)

  output_buffer_size = batch_size * 1000

  source_eos_id = tf.to_int32(source_vocab.lookup(tf.constant(eos)))
  target_sos_id = tf.to_int32(target_vocab.lookup(tf.constant(sos)))
  target_eos_id = tf.to_int32(target_vocab.lookup(tf.constant(eos)))

  dataset = tf.data.Dataset.zip((source_dataset, target_dataset))
  dataset = dataset.map(
    lambda src, tgt: (tf.string_split([src]).values,
                      tf.string_split([tgt]).values)).prefetch(output_buffer_size)
  dataset = dataset.filter(
    lambda src, tgt: tf.logical_and(tf.size(src) > 0, tf.size(tgt) > 0))
  dataset = dataset.map(
    lambda src, tgt: (src[:source_max_length], tgt[:target_max_length]))
  dataset = dataset.prefetch(output_buffer_size)

  dataset = dataset.map(
    lambda src, tgt: (tf.to_int32(source_vocab.lookup(src)),
                      tf.to_int32(target_vocab.lookup(tgt))))
  dataset = dataset.prefetch(output_buffer_size)

  dataset = dataset.map(
    lambda src, tgt: (tf.reverse(src, axis=[0]),
                      tf.concat(([target_sos_id], tgt), 0),
                      tf.concat((tgt, [target_eos_id]), 0))).prefetch(output_buffer_size)

  dataset = dataset.map(
    lambda src, tgt_in, tgt_out: (
      src, tgt_in, tgt_out, tf.size(src), tf.size(tgt_in))).prefetch(output_buffer_size)

  dataset = dataset.shuffle(100).repeat().padded_batch(
    batch_size,
    padded_shapes=(tf.TensorShape([None]),
                   tf.TensorShape([None]),
                   tf.TensorShape([None]),
                   tf.TensorShape([]),
                   tf.TensorShape([])),
    padding_values=(source_eos_id,
                    target_eos_id,
                    target_eos_id,
                    0,
                    0))

  iterator = dataset.make_initializable_iterator()

  return iterator.get_next(), iterator.initializer, source_vocab, target_vocab
//  [python]
def create_input_data(source_data_file, target_data_file,
                      source_vocab_file, target_vocab_file,
                      batch_size, unk_id, sos, eos,
                      source_max_length, target_max_length):
  source_dataset = tf.data.TextLineDataset(tf.gfile.Glob(source_data_file))
  target_dataset = tf.data.TextLineDataset(tf.gfile.Glob(target_data_file))
  source_vocab = lookup_ops.index_table_from_file(
    source_vocab_file, default_value=unk_id)
  target_vocab = lookup_ops.index_table_from_file(
    target_vocab_file, default_value=unk_id)

  output_buffer_size = batch_size * 1000

  source_eos_id = tf.cast(source_vocab.lookup(tf.constant(eos)), tf.int32)
  target_sos_id = tf.cast(target_vocab.lookup(tf.constant(sos)), tf.int32)
  target_eos_id = tf.cast(target_vocab.lookup(tf.constant(eos)), tf.int32)

  dataset = tf.data.Dataset.zip((source_dataset, target_dataset))
  dataset = dataset.map(
    lambda src, tgt: (tf.string_split([src]).values,
                      tf.string_split([tgt]).values)).prefetch(output_buffer_size)
  dataset = dataset.filter(
    lambda src, tgt: tf.logical_and(tf.size(src) > 0, tf.size(tgt) > 0))
  dataset = dataset.map(
    lambda src, tgt: (src[:source_max_length], tgt[:target_max_length]))
  dataset = dataset.prefetch(output_buffer_size)

  dataset = dataset.map(
    lambda src, tgt: (tf.cast(source_vocab.lookup(src), tf.int32),
                      tf.cast(target_vocab.lookup(tgt), tf.int32)))
  dataset = dataset.prefetch(output_buffer_size)

  dataset = dataset.map(
    lambda src, tgt: (tf.reverse(src, axis=[0]),
                      tf.concat(([target_sos_id], tgt), 0),
                      tf.concat((tgt, [target_eos_id]), 0))).prefetch(output_buffer_size)

  dataset = dataset.map(
    lambda src, tgt_in, tgt_out: (
      src, tgt_in, tgt_out, tf.size(src), tf.size(tgt_in))).prefetch(output_buffer_size)

  dataset = dataset.shuffle(100).repeat().padded_batch(
    batch_size,
    padded_shapes=(tf.TensorShape([None]),
                   tf.TensorShape([None]),
                   tf.TensorShape([None]),
                   tf.TensorShape([]),
                   tf.TensorShape([])),
    padding_values=(source_eos_id,
                    target_eos_id,
                    target_eos_id,
                    0,
                    0))

  iterator = dataset.make_initializable_iterator()

  return iterator.get_next(), iterator.initializer, source_vocab, target_vocab
//  [python]
def prototypical_clustering_layer(nclasses,
                                  x_train,
                                  y_train,
                                  x_test,
                                  phi,
                                  num_cluster_steps,
                                  lambd=0.1,
                                  alpha=0.0):
  """Computes the prototypes, cluster centers, with additional clustering on
    the validation data.
  Args:
    x_train: [N, D], Train data.
    y_train: [N], Train class labels.
    x_test: [N, D], Test data.
    phi: Feature extractor function.
  Returns:
    logits: [N, K], Test prediction.
  """
  protos = [None] * nclasses
  x_all = concat([x_train, x_test], 0)
  h, _ = phi(x_all, reuse=None, is_training=True)
  num_x_train = tf.shape(x_train)[0]
  h_train = h[:num_x_train, :]
  h_test = h[num_x_train:, :]

  # Initialize cluster center.
  for kk in range(nclasses):
    ksel = tf.expand_dims(tf.to_float(tf.equal(y_train, kk)),
                          1)  # [N, 1]
    protos[kk] = tf.reduce_sum(h_train * ksel, [0], keep_dims=True)  # [N, D]
    protos[kk] /= tf.reduce_sum(ksel)
  protos = concat(protos, 0)  # [K, D]

  # Run clustering.
  for tt in range(num_cluster_steps):
    all_data = concat([h_train, h_test], 0)
    # Label assignment.
    prob = assign_cluster(protos, all_data)
    # Prototype update.
    ## This is a vanilla version.
    ## Probably want to impose a constraint that each training example can only
    ## be in one cluster.
    protos = update_cluster(all_data, prob)

  # Be cautious here!!
  # Returns the logits and unsupervised clustering loss.
  uloss = sq_dist_loss(protos, all_data)
  logits = compute_logits(protos, h_test)
  return logits, uloss
//  [python]
def prototypical_clustering_layer(nclasses,
                                  x_train,
                                  y_train,
                                  x_test,
                                  phi,
                                  num_cluster_steps,
                                  lambd=0.1,
                                  alpha=0.0):
  """Computes the prototypes, cluster centers, with additional clustering on
    the validation data.
  Args:
    x_train: [N, D], Train data.
    y_train: [N], Train class labels.
    x_test: [N, D], Test data.
    phi: Feature extractor function.
  Returns:
    logits: [N, K], Test prediction.
  """
  protos = [None] * nclasses
  x_all = concat([x_train, x_test], 0)
  h, _ = phi(x_all, reuse=None, is_training=True)
  num_x_train = tf.shape(x_train)[0]
  h_train = h[:num_x_train, :]
  h_test = h[num_x_train:, :]

  # Initialize cluster center.
  for kk in range(nclasses):
    ksel = tf.expand_dims(tf.cast(tf.equal(y_train, kk), h_train.dtype),
                          1)  # [N, 1]
    protos[kk] = tf.reduce_sum(h_train * ksel, [0], keep_dims=True)  # [N, D]
    protos[kk] /= tf.reduce_sum(ksel)
  protos = concat(protos, 0)  # [K, D]

  # Run clustering.
  for tt in range(num_cluster_steps):
    all_data = concat([h_train, h_test], 0)
    # Label assignment.
    prob = assign_cluster(protos, all_data)
    # Prototype update.
    ## This is a vanilla version.
    ## Probably want to impose a constraint that each training example can only
    ## be in one cluster.
    protos = update_cluster(all_data, prob)

  # Be cautious here!!
  # Returns the logits and unsupervised clustering loss.
  uloss = sq_dist_loss(protos, all_data)
  logits = compute_logits(protos, h_test)
  return logits, uloss
//  [python]
def _model_fn_train(
      self,
      mode,
      total_loss,
      batches_per_epoch,
      num_epochs_per_decay,
      initial_learning_rate,
      learning_rate_decay_factor,
      rmsprop_decay,
      rmsprop_momentum,
      rmsprop_epsilon,
      moving_average_decay,
  ):
    """This is the TRAIN part of model_fn."""
    if mode != tf_estimator.ModeKeys.TRAIN:
      return None, None

    # Configure the learning rate using an exponetial decay.
    global_step = tf.compat.v1.train.get_or_create_global_step()
    current_epoch = tf.cast(global_step, tf.float32) / batches_per_epoch
    decay_steps = int(1.0 * batches_per_epoch * num_epochs_per_decay)

    learning_rate = tf.compat.v1.train.exponential_decay(
        learning_rate=initial_learning_rate,
        global_step=global_step,
        decay_steps=decay_steps,
        decay_rate=learning_rate_decay_factor,
        staircase=True,
    )

    # Set a minimum boundary for the learning rate to be a fixed value of 1e-9.
    # It's common to see these tf.max(...) operations when training inception,
    # with a max of 1e-4 * initial_learning_rate but this makes it hard to
    # explore learning rate schedules that decay quickly or by a lot of each
    # step. Here we just use a very small constant 1e-9 as the minimum value.
    learning_rate = tf.maximum(learning_rate, 1e-9, name='learning_rate')

    optimizer = tf.compat.v1.train.RMSPropOptimizer(
        learning_rate,
        rmsprop_decay,
        momentum=rmsprop_momentum,
        epsilon=rmsprop_epsilon,
    )
    if self.use_tpu:
      optimizer = tpu_optimizer.CrossShardOptimizer(optimizer)
    update_ops = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS)
    with tf.control_dependencies(update_ops):
      train_op = optimizer.minimize(total_loss, global_step=global_step)

    # NB. In the inception code this was "tf.trainable_variables()
    # + tf.moving_average_variables()", but we've settled on just
    # tf.model_variables() in the existing production DV2.
    variables_to_average = tf.compat.v1.model_variables()
    variable_averages = tf.train.ExponentialMovingAverage(
        decay=moving_average_decay, num_updates=global_step
    )
    with tf.control_dependencies([train_op]), tf.compat.v1.name_scope(
        'moving_average'
    ):
      train_op = variable_averages.apply(variables_to_average)
    tf.compat.v1.add_to_collection(tf.compat.v1.GraphKeys.UPDATE_OPS, train_op)

    # Compute the current epoch and associated learning rate from global_step.
    metric_dict = {
        'current_epoch': current_epoch,
        'total_loss': total_loss,
        'learning_rate': learning_rate,
    }
    host_call = self.construct_scalar_host_call(
        metric_dict=metric_dict, model_dir=self.model_dir, prefix='training/'
    )

    return train_op, host_call
//  [python]
def prototypical_clustering_gmm_layer(nclasses,
                                      x_train,
                                      y_train,
                                      x_test,
                                      phi,
                                      num_cluster_steps,
                                      lambd=0.1,
                                      alpha=0.0):
  """Computes the prototypes, cluster centers, with additional clustering on
    the validation data.
  Args:
    x_train: [N, D], Train data.
    y_train: [N], Train class labels.
    x_test: [N, D], Test data.
    phi: Feature extractor function.
  Returns:
    logits: [N, K], Test prediction.
  """
  protos = [None] * nclasses
  covar = [None] * nclasses
  x_all = concat([x_train, x_test], 0)
  h, _ = phi(x_all, reuse=None, is_training=True)
  num_x_train = tf.shape(x_train)[0]
  h_train = h[:num_x_train, :]
  h_test = h[num_x_train:, :]
  ndim = tf.shape(h)[1]

  # Initialize cluster center.
  for kk in range(nclasses):
    ksel = tf.expand_dims(tf.to_float(tf.equal(y_train, kk)),
                          1)  # [N, 1]
    protos[kk] = tf.reduce_sum(h_train * ksel, [0], keep_dims=True)  # [N, D]
    protos[kk] /= tf.reduce_sum(ksel)
    #covar[kk] = tf.expand_dims(tf.eye(ndim), 0)
    covar[kk] = tf.ones([1, ndim])  # diagonal
  protos = concat(protos, 0)  # [K, D]
  covar = concat(covar, 0)  # [K, D, D]

  # Run clustering.
  for tt in range(num_cluster_steps):
    all_data = concat([h_train, h_test], 0)
    # Label assignment.
    prob = assign_gmm_diag_cluster(protos, covar, all_data, nclasses)
    protos, covar = update_gmm_diag_cluster(all_data, prob)

  # Be cautious here!!
  uloss = mh_diag_dist_loss(protos, covar, all_data, nclasses)
  logits = compute_gmm_diag_logits(protos, covar, h_test, nclasses)
  return logits, uloss
//  [python]
def prototypical_clustering_gmm_layer(nclasses,
                                      x_train,
                                      y_train,
                                      x_test,
                                      phi,
                                      num_cluster_steps,
                                      lambd=0.1,
                                      alpha=0.0):
  """Computes the prototypes, cluster centers, with additional clustering on
    the validation data.
  Args:
    x_train: [N, D], Train data.
    y_train: [N], Train class labels.
    x_test: [N, D], Test data.
    phi: Feature extractor function.
  Returns:
    logits: [N, K], Test prediction.
  """
  protos = [None] * nclasses
  covar = [None] * nclasses
  x_all = concat([x_train, x_test], 0)
  h, _ = phi(x_all, reuse=None, is_training=True)
  num_x_train = tf.shape(x_train)[0]
  h_train = h[:num_x_train, :]
  h_test = h[num_x_train:, :]
  ndim = tf.shape(h)[1]

  # Initialize cluster center.
  for kk in range(nclasses):
    ksel = tf.expand_dims(tf.cast(tf.equal(y_train, kk), h_train.dtype),
                          1)  # [N, 1]
    protos[kk] = tf.reduce_sum(h_train * ksel, [0], keep_dims=True)  # [N, D]
    protos[kk] /= tf.reduce_sum(ksel)
    #covar[kk] = tf.expand_dims(tf.eye(ndim), 0)
    covar[kk] = tf.ones([1, ndim])  # diagonal
  protos = concat(protos, 0)  # [K, D]
  covar = concat(covar, 0)  # [K, D, D]

  # Run clustering.
  for tt in range(num_cluster_steps):
    all_data = concat([h_train, h_test], 0)
    # Label assignment.
    prob = assign_gmm_diag_cluster(protos, covar, all_data, nclasses)
    protos, covar = update_gmm_diag_cluster(all_data, prob)

  # Be cautious here!!
  uloss = mh_diag_dist_loss(protos, covar, all_data, nclasses)
  logits = compute_gmm_diag_logits(protos, covar, h_test, nclasses)
  return logits, uloss
//  [python]
def _optimal_path(self, cm_, tb_):
        r"""
        Compute optimal path (state at each time t) given tensors cm_ & tb_
        of shapes (None, Ns, T). Output is of shape (None, T)
        cm_: cumulative metrics for each state at time t(0 to T)
        tb_: traceback state for each state at time t(0 to T)
        """
        # tb and ca are of shape (batch x self._ns x num_syms)
        assert(tb_.get_shape()[1] == self._ns), "Invalid shape."
        optst_ta = tf.TensorArray(tf.int32, size=tb_.shape[-1],
                                  dynamic_size=False,
                                  clear_after_read=False)
        if self._terminate:
            opt_term_state = tf.zeros((tf.shape(cm_)[0],), tf.int32)
        else:
            opt_term_state =tf.to_int32(tf.argmin(cm_[:, :, -1], axis=1))
        optst_ta = optst_ta.write(tb_.shape[-1]-1,opt_term_state)

        for sym in tf.range(tb_.shape[-1]-1, 0, -1):
            opt_st = optst_ta.read(sym)[:,None]

            idx_ = tf.concat([tf.range(tf.shape(cm_)[0])[:,None], opt_st],
                             axis=1)
            opt_st_tminus1 = tf.gather_nd(tb_[:, :, sym], idx_)

            optst_ta = optst_ta.write(sym-1, opt_st_tminus1)

        return tf.transpose(optst_ta.stack())
//  [python]
def _optimal_path(self, cm_, tb_):
        r"""
        Compute optimal path (state at each time t) given tensors cm_ & tb_
        of shapes (None, Ns, T). Output is of shape (None, T)
        cm_: cumulative metrics for each state at time t(0 to T)
        tb_: traceback state for each state at time t(0 to T)
        """
        # tb and ca are of shape (batch x self._ns x num_syms)
        assert(tb_.get_shape()[1] == self._ns), "Invalid shape."
        optst_ta = tf.TensorArray(tf.int32, size=tb_.shape[-1],
                                  dynamic_size=False,
                                  clear_after_read=False)
        if self._terminate:
            opt_term_state = tf.zeros((tf.shape(cm_)[0],), tf.int32)
        else:
            opt_term_state =tf.cast(tf.argmin(cm_[:, :, -1], axis=1), tf.int32)
        optst_ta = optst_ta.write(tb_.shape[-1]-1,opt_term_state)

        for sym in tf.range(tb_.shape[-1]-1, 0, -1):
            opt_st = optst_ta.read(sym)[:,None]

            idx_ = tf.concat([tf.range(tf.shape(cm_)[0])[:,None], opt_st],
                             axis=1)
            opt_st_tminus1 = tf.gather_nd(tb_[:, :, sym], idx_)

            optst_ta = optst_ta.write(sym-1, opt_st_tminus1)

        return tf.transpose(optst_ta.stack())
//  [python]
def _op_bits_path(self, paths):
        r"""
        Given a path, compute the input bit stream that results in the path.
        Used in call() where the input is optimal path (seq of states) such
        as the path returned by _return_optimal.
        """
        paths = t.cast(paths, t.int32)
        ip_bits = tf.TensorArray(tf.int32,
                                  size=paths.shape[-1]-1,
                                  dynamic_size=False,
                                  clear_after_read=False)
        dec_syms = tf.TensorArray(tf.int32,
                                  size=paths.shape[-1]-1,
                                  dynamic_size=False,
                                  clear_after_read=False)
        ni = self._trellis.ni
        ip_sym_mask = tf.range(ni)[None, :]

        for sym in tf.range(1, paths.shape[-1]):

            # gather index from paths to enable XLA
            # replaces p_idx = paths[:,sym-1:sym+1]
            p_idx = tf.gather(paths, [sym-1, sym], axis=-1)
            dec_ = tf.gather_nd(self._trellis.op_mat, p_idx)

            dec_syms = dec_syms.write(sym-1, value=dec_)
            # bs x ni boolean tensor. Each row has a True and False. True
            # corresponds to input_bit which produced the next state (t=sym)
            match_st = tf.math.equal(
                tf.gather(self._trellis.to_nodes,paths[:, sym-1]),
                tf.tile(paths[:, sym][:, None], [1, 2])
                )

            # tf.boolean_mask throws error in XLA mode
            #ip_bit = tf.boolean_mask(ip_sym_mask, match_st)

            ip_bit_ = tf.where(match_st,
                               ip_sym_mask,
                               tf.zeros_like(ip_sym_mask))
            ip_bit = tf.reduce_sum(ip_bit_, axis=-1)
            ip_bits = ip_bits.write(sym-1, ip_bit)

        ip_bit_vec_est = tf.transpose(ip_bits.stack())
        ip_sym_vec_est = tf.transpose(dec_syms.stack())

        return ip_bit_vec_est, ip_sym_vec_est
//  [python]
def _op_bits_path(self, paths):
        r"""
        Given a path, compute the input bit stream that results in the path.
        Used in call() where the input is optimal path (seq of states) such
        as the path returned by _return_optimal.
        """
        paths = tf.cast(paths, tf.int32)
        ip_bits = tf.TensorArray(tf.int32,
                                  size=paths.shape[-1]-1,
                                  dynamic_size=False,
                                  clear_after_read=False)
        dec_syms = tf.TensorArray(tf.int32,
                                  size=paths.shape[-1]-1,
                                  dynamic_size=False,
                                  clear_after_read=False)
        ni = self._trellis.ni
        ip_sym_mask = tf.range(ni)[None, :]

        for sym in tf.range(1, paths.shape[-1]):

            # gather index from paths to enable XLA
            # replaces p_idx = paths[:,sym-1:sym+1]
            p_idx = tf.gather(paths, [sym-1, sym], axis=-1)
            dec_ = tf.gather_nd(self._trellis.op_mat, p_idx)

            dec_syms = dec_syms.write(sym-1, value=dec_)
            # bs x ni boolean tensor. Each row has a True and False. True
            # corresponds to input_bit which produced the next state (t=sym)
            match_st = tf.math.equal(
                tf.gather(self._trellis.to_nodes,paths[:, sym-1]),
                tf.tile(paths[:, sym][:, None], [1, 2])
                )

            # tf.boolean_mask throws error in XLA mode
            #ip_bit = tf.boolean_mask(ip_sym_mask, match_st)

            ip_bit_ = tf.where(match_st,
                               ip_sym_mask,
                               tf.zeros_like(ip_sym_mask))
            ip_bit = tf.reduce_sum(ip_bit_, axis=-1)
            ip_bits = ip_bits.write(sym-1, ip_bit)

        ip_bit_vec_est = tf.transpose(ip_bits.stack())
        ip_sym_vec_est = tf.transpose(dec_syms.stack())

        return ip_bit_vec_est, ip_sym_vec_est