jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def focal_loss_softmax(labels, logits, gamma=2):
    """
    https://github.com/fudannlp16/focal-loss/blob/master/focal_loss.py
    Computer focal loss for multi classification
    Args:
      labels: A int32 tensor of shape [batch_size].
      logits: A float32 tensor of shape [batch_size,num_classes].
      gamma: A scalar for focal loss gamma hyper-parameter.
    Returns:
      A tensor of the same shape as `lables`
    """
    y_pred=tf.nn.softmax(logits,axis=-1) # [batch_size,num_classes]
    cross_ent = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=tf.squeeze(labels, axis=1), logits=logits)
    cross_ent = tf.expand_dims(cross_ent, 1)
    labels = tf.cast(labels, tf.int32)
    L=((1.0-tf.batch_gather(y_pred, labels))**gamma)*cross_ent
    return L
//  [python]
def call(self, input_ids, past_key_values_length=0):
        past_key_values_length = t.cast(past_key_values_length, t.int32)
        gather_indices = tf.range(tf.shape(input_ids)[1], delta=1) + past_key_values_length
        return tf.gather(self.weight, gather_indices)
//  [python]
def call(self, input_ids, past_key_values_length=0):
        past_key_values_length = tf.cast(past_key_values_length, tf.int32)
        gather_indices = tf.range(tf.shape(input_ids)[1], delta=1) + past_key_values_length
        return tf.gather(self.weight, gather_indices)
//  [python]
def sinusoidal_embedding_init(shape, dtype=tf.float32) -> tf.Tensor:
    """Returns sinusoids for positional embedding"""
    length, channels = shape
    if channels % 2 != 0:
        raise ValueError(
            f"Number of channels has to be divisible by 2 for sinusoidal positional embeddings, got {channels} channels."
        )
    log_timescale_increment = math.log(10000) / (channels // 2 - 1)
    inv_timescales = tf.exp(-log_timescale_increment * tf.range(channels // 2, dtype=tf.float32))
    scaled_time = tf.reshape(tf.range(length, dtype=tf.float32), (-1, 1)) * tf.reshape(inv_timescales, (1, -1))
    return tf.to_float(tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=1))
//  [python]
def sinusoidal_embedding_init(shape, dtype=tf.float32) -> tf.Tensor:
    """Returns sinusoids for positional embedding"""
    length, channels = shape
    if channels % 2 != 0:
        raise ValueError(
            f"Number of channels has to be divisible by 2 for sinusoidal positional embeddings, got {channels} channels."
        )
    log_timescale_increment = math.log(10000) / (channels // 2 - 1)
    inv_timescales = tf.exp(-log_timescale_increment * tf.range(channels // 2, dtype=tf.float32))
    scaled_time = tf.reshape(tf.range(length, dtype=tf.float32), (-1, 1)) * tf.reshape(inv_timescales, (1, -1))
    return tf.cast(tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=1), dtype)
//  [python]
def test_classification_loss_scalar_weight_scales_inputs(self):
    n = 12
    num_classes = 5
    weight = 4.0
    logits = tf.random.uniform((n, num_classes), maxval=1)
    labels = tf.cast(
        tf.math.round(tf.random.uniform(
            (n, 1), minval=1, maxval=num_classes)) - 1,
        dtype=tf.int32)
    outputs = {
        standard_fields.DetectionResultFields.object_semantic_voxels: logits
    }
    inputs1 = {
        standard_fields.InputDataFields.object_class_voxels: labels,
        standard_fields.InputDataFields.voxel_loss_weights: 1.0,
    }
    inputs2 = {
        standard_fields.InputDataFields.object_class_voxels: labels,
        standard_fields.InputDataFields.voxel_loss_weights: weight,
    }
    loss_unscaled = classification_losses.classification_loss(inputs1, outputs)
    loss_scaled = classification_losses.classification_loss(inputs2, outputs)
    self.assertEqual(loss_unscaled.numpy() * weight, loss_scaled.numpy())
//  [python]
def test_classification_loss_scalar_weight_scales_inputs_batch(self):
    batch_size = 4
    n = 12
    num_classes = 5
    weight = 4.0
    logits = tf.random.uniform((batch_size, n, num_classes), maxval=1)
    labels = tf.to_int32(
        tf.math.round(
            tf.random.uniform(
                (batch_size, n, 1), minval=1, maxval=num_classes)) - 1)
    outputs = {
        standard_fields.DetectionResultFields.object_semantic_voxels: logits
    }
    inputs1 = {
        standard_fields.InputDataFields.object_class_voxels:
            labels,
        standard_fields.InputDataFields.voxel_loss_weights:
            1.0,
        standard_fields.InputDataFields.num_valid_voxels:
            tf.ones([batch_size], dtype=tf.int32) * n,
    }
    inputs2 = {
        standard_fields.InputDataFields.object_class_voxels:
            labels,
        standard_fields.InputDataFields.voxel_loss_weights:
            weight,
        standard_fields.InputDataFields.num_valid_voxels:
            tf.ones([batch_size], dtype=tf.int32) * n,
    }
    loss_unscaled = classification_losses.classification_loss(inputs1, outputs)
    loss_scaled = classification_losses.classification_loss(inputs2, outputs)
    self.assertEqual(loss_unscaled.numpy() * weight, loss_scaled.numpy())
//  [python]
def test_classification_loss_scalar_weight_scales_inputs_batch(self):
    batch_size = 4
    n = 12
    num_classes = 5
    weight = 4.0
    logits = tf.random.uniform((batch_size, n, num_classes), maxval=1)
    labels = tf.cast(
        tf.math.round(
            tf.random.uniform(
                (batch_size, n, 1), minval=1, maxval=num_classes)) - 1,
        dtype=tf.int32)
    outputs = {
        standard_fields.DetectionResultFields.object_semantic_voxels: logits
    }
    inputs1 = {
        standard_fields.InputDataFields.object_class_voxels:
            labels,
        standard_fields.InputDataFields.voxel_loss_weights:
            1.0,
        standard_fields.InputDataFields.num_valid_voxels:
            tf.ones([batch_size], dtype=tf.int32) * n,
    }
    inputs2 = {
        standard_fields.InputDataFields.object_class_voxels:
            labels,
        standard_fields.InputDataFields.voxel_loss_weights:
            weight,
        standard_fields.InputDataFields.num_valid_voxels:
            tf.ones([batch_size], dtype=tf.int32) * n,
    }
    loss_unscaled = classification_losses.classification_loss(inputs1, outputs)
    loss_scaled = classification_losses.classification_loss(inputs2, outputs)
    self.assertEqual(loss_unscaled.numpy() * weight, loss_scaled.numpy())
//  [python]
def call(self, inputs):
        K = self.K
        outputs = []
        for count in range(self.n_tasks):
            # Similarity values
            similarity = inputs[:, 2 * K * count:(2 * K * count + K)]
            # Labels for all top K similar samples
            ys = t.cast(inputs[:, (2 * K * count + K):2 * K * (count + 1)],
                         t.int32)

            R = self.b + self.W[0] * similarity + self.W[1] * tf.constant(
                np.arange(K) + 1, dtype=tf.float32)
            R = tf.sigmoid(R)
            z = tf.reduce_sum(R * tf.gather(self.V, ys), axis=1) + self.b2
            outputs.append(tf.reshape(z, shape=[-1, 1]))
        loss = (tf.nn.l2_loss(self.W) + tf.nn.l2_loss(self.V) +
                tf.nn.l2_loss(self.b) + tf.nn.l2_loss(self.b2)) * self.penalty
        self.add_loss(loss)
        return tf.concat(outputs, axis=1)
//  [python]
def call(self, inputs):
        K = self.K
        outputs = []
        for count in range(self.n_tasks):
            # Similarity values
            similarity = inputs[:, 2 * K * count:(2 * K * count + K)]
            # Labels for all top K similar samples
            ys = tf.cast(inputs[:, (2 * K * count + K):2 * K * (count + 1)],
                         tf.int32)

            R = self.b + self.W[0] * similarity + self.W[1] * tf.constant(
                np.arange(K) + 1, dtype=tf.float32)
            R = tf.sigmoid(R)
            z = tf.reduce_sum(R * tf.gather(self.V, ys), axis=1) + self.b2
            outputs.append(tf.reshape(z, shape=[-1, 1]))
        loss = (tf.nn.l2_loss(self.W) + tf.nn.l2_loss(self.V) +
                tf.nn.l2_loss(self.b) + tf.nn.l2_loss(self.b2)) * self.penalty
        self.add_loss(loss)
        return tf.concat(outputs, axis=1)
//  [python]
def test_classification_loss_using_mask_iou_func(self):
    embeddings = tf.constant([[[1.0, 0.0, 0.0],
                               [0.2, 0.2, 0.2],
                               [0.0, 1.0, 0.0],
                               [0.0, 0.0, 1.0],
                               [0.5, 0.5, 0.0],
                               [1.0, 0.0, 0.0],
                               [0.0, 1.0, 0.0],
                               [0.0, 0.0, 1.0]]], dtype=tf.float32)
    instance_ids = tf.constant([[0, 1, 2, 3, 1, 0, 0, 0]], dtype=tf.int32)
    class_labels = tf.constant([[[0], [1], [1], [1], [1], [0], [2], [2]]],
                               dtype=tf.int32)
    logits = tf.random.uniform([1, 8, 5],
                               minval=-1.0,
                               maxval=1.0,
                               dtype=tf.float32)
    valid_mask = tf.cast(
        tf.constant([[1, 1, 1, 0, 0, 1, 0, 0]], dtype=tf.int32), dtype=tf.bool)
    loss = classification_losses.classification_loss_using_mask_iou_func(
        embeddings=embeddings,
        logits=logits,
        instance_ids=instance_ids,
        class_labels=class_labels,
        num_samples=4,
        valid_mask=valid_mask,
        max_instance_id=5,
        similarity_strategy='dotproduct',
        is_balanced=True)
    self.assertGreater(loss.numpy(), 0.0)
//  [python]
def build_model_with_precision(pp, mm, ii, tt, *args, **kwargs):
  """Build model with its inputs/params for a specified precision context.

  This is highly specific to this codebase, and not intended to be general API.
  Advanced users only. DO NOT use it if you don't know what it does.
  NOTE: short argument names are intended to avoid conficts with kwargs.

  Args:
    pp: A string, precision policy name, such as "mixed_float16".
    mm: A function, for rmodel builder.
    ii: A tensor, for model inputs.
    tt: A bool, If true, it is for training; otherwise, it is for eval.
    *args: A list of model arguments.
    **kwargs: A dict, extra model parameters.

  Returns:
    the output of mm model.
  """
  del tt
  if pp == 'mixed_bfloat16':
    set_precision_policy(pp)
    inputs = t.cast(ii, tf.bfloat16)
    with tf.compat.v1.tpu.bfloat16_scope():
      outputs = mm(inputs, *args, **kwargs)
    set_precision_policy('float32')
  elif pp == 'mixed_float16':
    set_precision_policy(pp)
    inputs = t.cast(ii, tf.float16)
    with float16_scope():
      outputs = mm(inputs, *args, **kwargs)
    set_precision_policy('float32')
  elif not pp or pp == 'float32':
    outputs = mm(ii, *args, **kwargs)
  else:
    raise ValueError('Unknow precision name {}'.format(pp))

  # Users are responsible to convert the dtype of all outputs.
  return outputs
//  [python]
def build_model_with_precision(pp, mm, ii, tt, *args, **kwargs):
  """Build model with its inputs/params for a specified precision context.

  This is highly specific to this codebase, and not intended to be general API.
  Advanced users only. DO NOT use it if you don't know what it does.
  NOTE: short argument names are intended to avoid conficts with kwargs.

  Args:
    pp: A string, precision policy name, such as "mixed_float16".
    mm: A function, for rmodel builder.
    ii: A tensor, for model inputs.
    tt: A bool, If true, it is for training; otherwise, it is for eval.
    *args: A list of model arguments.
    **kwargs: A dict, extra model parameters.

  Returns:
    the output of mm model.
  """
  del tt
  if pp == 'mixed_bfloat16':
    set_precision_policy(pp)
    inputs = tf.cast(ii, tf.bfloat16)
    with tf.compat.v1.tpu.bfloat16_scope():
      outputs = mm(inputs, *args, **kwargs)
    set_precision_policy('float32')
  elif pp == 'mixed_float16':
    set_precision_policy(pp)
    inputs = tf.cast(ii, tf.float16)
    with float16_scope():
      outputs = mm(inputs, *args, **kwargs)
    set_precision_policy('float32')
  elif not pp or pp == 'float32':
    outputs = mm(ii, *args, **kwargs)
  else:
    raise ValueError('Unknow precision name {}'.format(pp))

  # Users are responsible to convert the dtype of all outputs.
  return outputs
//  [python]
def cast(xs, dtype, whiltelist=(tf.int32, tf.int64, tf.bool)):
    def func(x):
        if x.dtype != dtype and all(x.dtype != wdtype for wdtype in whiltelist):
            x = t.cast(x, dtype)
        return x
    return tf.nest.map_structure(func, xs)
//  [python]
def cast(xs, dtype, whiltelist=(tf.int32, tf.int64, tf.bool)):
    def func(x):
        if x.dtype != dtype and all(x.dtype != wdtype for wdtype in whiltelist):
            x = tf.cast(x, dtype)
        return x
    return tf.nest.map_structure(func, xs)
//  [python]
def smooth_l1_loss(y_true, y_pred):
    '''Implements Smooth-L1 loss.
    
    Args
    ---
        y_true and y_pred are typically: [N, 4], but could be any shape.
    '''
    diff = tf.abs(y_true - y_pred)
    less_than_one = tf.to_float(tf.less(diff, 1.0))
    loss = (less_than_one * 0.5 * diff**2) + (1 - less_than_one) * (diff - 0.5)
    return loss
//  [python]
def smooth_l1_loss(y_true, y_pred):
    '''Implements Smooth-L1 loss.
    
    Args
    ---
        y_true and y_pred are typically: [N, 4], but could be any shape.
    '''
    diff = tf.abs(y_true - y_pred)
    less_than_one = tf.cast(tf.less(diff, 1.0), tf.float32)
    loss = (less_than_one * 0.5 * diff**2) + (1 - less_than_one) * (diff - 0.5)
    return loss
//  [python]
def _distribution(self, time_step, policy_state):
    if time_step.step_type.shape:
      if (
          tf.is_tensor(time_step.step_type)
          and time_step.step_type.shape.rank > 0
      ):
        batch_size = time_step.step_type.get_shape().as_list()[0]
      else:
        batch_size = 1
    else:
      batch_size = 1
    # Sample from the posterior distribution.
    posterior_dist = tfd.Beta(self._alpha, self._beta)
    predicted_reward_sampled = posterior_dist.sample([batch_size])
    predicted_reward_means_1d = tf.stack(
        [
            self._alpha[k] / (self._alpha[k] + self._beta[k])
            for k in range(self._expected_num_actions)
        ],
        axis=-1,
    )
    predicted_reward_means = tf.stack(
        [predicted_reward_means_1d for k in range(batch_size)], axis=0
    )

    mask = None
    if self._observation_and_action_constraint_splitter is not None:
      _, mask = self._observation_and_action_constraint_splitter(
          time_step.observation
      )

    # Argmax.
    if mask is not None:
      actions = policy_utilities.masked_argmax(
          predicted_reward_sampled, mask, output_type=self.action_spec.dtype
      )
    else:
      actions = tf.argmax(
          predicted_reward_sampled, axis=-1, output_type=self.action_spec.dtype
      )

    policy_info = policy_utilities.populate_policy_info(
        arm_observations=(),
        chosen_actions=actions,
        rewards_for_argmax=tf.to_float(predicted_reward_sampled),
        est_rewards=tf.to_float(predicted_reward_means),
        emit_policy_info=self._emit_policy_info,
        accepts_per_arm_features=False,
    )
    if policy_utilities.InfoFields.LOG_PROBABILITY in self._emit_policy_info:
      policy_info._replace(log_probability=tf.zeros([batch_size], tf.float32))

    return policy_step.PolicyStep(
        tfp.distributions.Deterministic(loc=actions), policy_state, policy_info
    )
//  [python]
def _distribution(self, time_step, policy_state):
    if time_step.step_type.shape:
      if (
          tf.is_tensor(time_step.step_type)
          and time_step.step_type.shape.rank > 0
      ):
        batch_size = time_step.step_type.get_shape().as_list()[0]
      else:
        batch_size = 1
    else:
      batch_size = 1
    # Sample from the posterior distribution.
    posterior_dist = tfd.Beta(self._alpha, self._beta)
    predicted_reward_sampled = posterior_dist.sample([batch_size])
    predicted_reward_means_1d = tf.stack(
        [
            self._alpha[k] / (self._alpha[k] + self._beta[k])
            for k in range(self._expected_num_actions)
        ],
        axis=-1,
    )
    predicted_reward_means = tf.stack(
        [predicted_reward_means_1d for k in range(batch_size)], axis=0
    )

    mask = None
    if self._observation_and_action_constraint_splitter is not None:
      _, mask = self._observation_and_action_constraint_splitter(
          time_step.observation
      )

    # Argmax.
    if mask is not None:
      actions = policy_utilities.masked_argmax(
          predicted_reward_sampled, mask, output_type=self.action_spec.dtype
      )
    else:
      actions = tf.argmax(
          predicted_reward_sampled, axis=-1, output_type=self.action_spec.dtype
      )

    policy_info = policy_utilities.populate_policy_info(
        arm_observations=(),
        chosen_actions=actions,
        rewards_for_argmax=tf.cast(predicted_reward_sampled, tf.float32),
        est_rewards=tf.cast(predicted_reward_means, tf.float32),
        emit_policy_info=self._emit_policy_info,
        accepts_per_arm_features=False,
    )
    if policy_utilities.InfoFields.LOG_PROBABILITY in self._emit_policy_info:
      policy_info._replace(log_probability=tf.zeros([batch_size], tf.float32))

    return policy_step.PolicyStep(
        tfp.distributions.Deterministic(loc=actions), policy_state, policy_info
    )
//  [python]
def rpn_class_loss(target_matchs, rpn_class_logits):
    '''RPN anchor classifier loss.
    
    Args
    ---
        target_matchs: [batch_size, num_anchors]. Anchor match type. 1=positive,
            -1=negative, 0=neutral anchor.
        rpn_class_logits: [batch_size, num_anchors, 2]. RPN classifier logits for FG/BG.
    '''

    # Get anchor classes. Convert the -1/+1 match to 0/1 values.
    anchor_class = tf.to_int32(tf.equal(target_matchs, 1))
    # Positive and Negative anchors contribute to the loss,
    # but neutral anchors (match value = 0) don't.
    indices = tf.where(tf.not_equal(target_matchs, 0))
    # Pick rows that contribute to the loss and filter out the rest.
    rpn_class_logits = tf.gather_nd(rpn_class_logits, indices)
    anchor_class = tf.gather_nd(anchor_class, indices)
    # Cross entropy loss
    # loss = tf.losses.sparse_softmax_cross_entropy(labels=anchor_class,
    #                                               logits=rpn_class_logits)

    num_classes = rpn_class_logits.shape[-1]
    # print(rpn_class_logits.shape)
    loss = keras.losses.categorical_crossentropy(tf.one_hot(anchor_class, depth=num_classes),
                                                 rpn_class_logits, from_logits=True)

    
    loss = tf.reduce_mean(loss) if tf.size(loss) > 0 else tf.constant(0.0)
    return loss
//  [python]
def rpn_class_loss(target_matchs, rpn_class_logits):
    '''RPN anchor classifier loss.
    
    Args
    ---
        target_matchs: [batch_size, num_anchors]. Anchor match type. 1=positive,
            -1=negative, 0=neutral anchor.
        rpn_class_logits: [batch_size, num_anchors, 2]. RPN classifier logits for FG/BG.
    '''

    # Get anchor classes. Convert the -1/+1 match to 0/1 values.
    anchor_class = tf.cast(tf.equal(target_matchs, 1), tf.int32)
    # Positive and Negative anchors contribute to the loss,
    # but neutral anchors (match value = 0) don't.
    indices = tf.where(tf.not_equal(target_matchs, 0))
    # Pick rows that contribute to the loss and filter out the rest.
    rpn_class_logits = tf.gather_nd(rpn_class_logits, indices)
    anchor_class = tf.gather_nd(anchor_class, indices)
    # Cross entropy loss
    # loss = tf.losses.sparse_softmax_cross_entropy(labels=anchor_class,
    #                                               logits=rpn_class_logits)

    num_classes = rpn_class_logits.shape[-1]
    # print(rpn_class_logits.shape)
    loss = keras.losses.categorical_crossentropy(tf.one_hot(anchor_class, depth=num_classes),
                                                 rpn_class_logits, from_logits=True)

    
    loss = tf.reduce_mean(loss) if tf.size(loss) > 0 else tf.constant(0.0)
    return loss
//  [python]
def discriminative_loss_single(prediction, correct_label, feature_dim, label_shape, 
                            delta_v, delta_d, param_var, param_dist, param_reg):
    
    ''' Discriminative loss for a single prediction/label pair.
    :param prediction: inference of network
    :param correct_label: instance label
    :feature_dim: feature dimension of prediction
    :param label_shape: shape of label
    :param delta_v: cutoff variance distance
    :param delta_d: curoff cluster distance
    :param param_var: weight for intra cluster variance
    :param param_dist: weight for inter cluster distances
    :param param_reg: weight regularization
    '''

    ### Reshape so pixels are aligned along a vector
    correct_label = tf.reshape(correct_label, [label_shape[1]*label_shape[0]])
    reshaped_pred = tf.reshape(prediction, [label_shape[1]*label_shape[0], feature_dim])

    ### Count instances
    unique_labels, unique_id, counts = tf.unique_with_counts(correct_label)
    counts = tf.to_float(counts)
    num_instances = tf.size(unique_labels)

    segmented_sum = tf.unsorted_segment_sum(reshaped_pred, unique_id, num_instances)

    mu = tf.div(segmented_sum, tf.reshape(counts, (-1, 1)))
    mu_expand = tf.gather(mu, unique_id)

    ### Calculate l_var
    distance = tf.norm(tf.subtract(mu_expand, reshaped_pred), axis=1)
    distance = tf.subtract(distance, delta_v)
    distance = tf.clip_by_value(distance, 0., distance)
    distance = tf.square(distance)

    l_var = tf.unsorted_segment_sum(distance, unique_id, num_instances)
    l_var = tf.div(l_var, counts)
    l_var = tf.reduce_sum(l_var)
    l_var = tf.divide(l_var, tf.to_float(num_instances))
    
    ### Calculate l_dist
    
    # Get distance for each pair of clusters like this:
    #   mu_1 - mu_1
    #   mu_2 - mu_1
    #   mu_3 - mu_1
    #   mu_1 - mu_2
    #   mu_2 - mu_2
    #   mu_3 - mu_2
    #   mu_1 - mu_3
    #   mu_2 - mu_3
    #   mu_3 - mu_3

    mu_interleaved_rep = tf.tile(mu, [num_instances, 1])
    mu_band_rep = tf.tile(mu, [1, num_instances])
    mu_band_rep = tf.reshape(mu_band_rep, (num_instances*num_instances, feature_dim))

    mu_diff = tf.subtract(mu_band_rep, mu_interleaved_rep)
    
    # Filter out zeros from same cluster subtraction
    intermediate_tensor = tf.reduce_sum(tf.abs(mu_diff),axis=1)
    zero_vector = tf.zeros(1, dtype=tf.float32)
    bool_mask = tf.not_equal(intermediate_tensor, zero_vector)
    mu_diff_bool = tf.boolean_mask(mu_diff, bool_mask)

    mu_norm = tf.norm(mu_diff_bool, axis=1)
    mu_norm = tf.subtract(2.*delta_d, mu_norm)
    mu_norm = tf.clip_by_value(mu_norm, 0., mu_norm)
    mu_norm = tf.square(mu_norm)

    l_dist = tf.reduce_mean(mu_norm)

    ### Calculate l_reg
    l_reg = tf.reduce_mean(tf.norm(mu, axis=1))

    param_scale = 1.
    l_var = param_var * l_var
    l_dist = param_dist * l_dist
    l_reg = param_reg * l_reg

    loss = param_scale*(l_var + l_dist + l_reg)
    
    return loss, l_var, l_dist, l_reg
//  [python]
def discriminative_loss_single(prediction, correct_label, feature_dim, label_shape, 
                            delta_v, delta_d, param_var, param_dist, param_reg):
    
    ''' Discriminative loss for a single prediction/label pair.
    :param prediction: inference of network
    :param correct_label: instance label
    :feature_dim: feature dimension of prediction
    :param label_shape: shape of label
    :param delta_v: cutoff variance distance
    :param delta_d: curoff cluster distance
    :param param_var: weight for intra cluster variance
    :param param_dist: weight for inter cluster distances
    :param param_reg: weight regularization
    '''

    ### Reshape so pixels are aligned along a vector
    correct_label = tf.reshape(correct_label, [label_shape[1]*label_shape[0]])
    reshaped_pred = tf.reshape(prediction, [label_shape[1]*label_shape[0], feature_dim])

    ### Count instances
    unique_labels, unique_id, counts = tf.unique_with_counts(correct_label)
    counts = tf.cast(counts, tf.float32)
    num_instances = tf.size(unique_labels)

    segmented_sum = tf.unsorted_segment_sum(reshaped_pred, unique_id, num_instances)

    mu = tf.div(segmented_sum, tf.reshape(counts, (-1, 1)))
    mu_expand = tf.gather(mu, unique_id)

    ### Calculate l_var
    distance = tf.norm(tf.subtract(mu_expand, reshaped_pred), axis=1)
    distance = tf.subtract(distance, delta_v)
    distance = tf.clip_by_value(distance, 0., distance)
    distance = tf.square(distance)

    l_var = tf.unsorted_segment_sum(distance, unique_id, num_instances)
    l_var = tf.div(l_var, counts)
    l_var = tf.reduce_sum(l_var)
    l_var = tf.divide(l_var, tf.cast(num_instances, tf.float32))
    
    ### Calculate l_dist
    
    # Get distance for each pair of clusters like this:
    #   mu_1 - mu_1
    #   mu_2 - mu_1
    #   mu_3 - mu_1
    #   mu_1 - mu_2
    #   mu_2 - mu_2
    #   mu_3 - mu_2
    #   mu_1 - mu_3
    #   mu_2 - mu_3
    #   mu_3 - mu_3

    mu_interleaved_rep = tf.tile(mu, [num_instances, 1])
    mu_band_rep = tf.tile(mu, [1, num_instances])
    mu_band_rep = tf.reshape(mu_band_rep, (num_instances*num_instances, feature_dim))

    mu_diff = tf.subtract(mu_band_rep, mu_interleaved_rep)
    
    # Filter out zeros from same cluster subtraction
    intermediate_tensor = tf.reduce_sum(tf.abs(mu_diff),axis=1)
    zero_vector = tf.zeros(1, dtype=tf.float32)
    bool_mask = tf.not_equal(intermediate_tensor, zero_vector)
    mu_diff_bool = tf.boolean_mask(mu_diff, bool_mask)

    mu_norm = tf.norm(mu_diff_bool, axis=1)
    mu_norm = tf.subtract(2.*delta_d, mu_norm)
    mu_norm = tf.clip_by_value(mu_norm, 0., mu_norm)
    mu_norm = tf.square(mu_norm)

    l_dist = tf.reduce_mean(mu_norm)

    ### Calculate l_reg
    l_reg = tf.reduce_mean(tf.norm(mu, axis=1))

    param_scale = 1.
    l_var = param_var * l_var
    l_dist = param_dist * l_dist
    l_reg = param_reg * l_reg

    loss = param_scale*(l_var + l_dist + l_reg)
    
    return loss, l_var, l_dist, l_reg
//  [python]
def rpn_bbox_loss(target_deltas, target_matchs, rpn_deltas):
    '''Return the RPN bounding box loss graph.
    
    Args
    ---
        target_deltas: [batch, num_rpn_deltas, (dy, dx, log(dh), log(dw))].
            Uses 0 padding to fill in unsed bbox deltas.
        target_matchs: [batch, anchors]. Anchor match type. 1=positive,
            -1=negative, 0=neutral anchor.
        rpn_deltas: [batch, anchors, (dy, dx, log(dh), log(dw))]
    '''
    def batch_pack(x, counts, num_rows):
        '''Picks different number of values from each row
        in x depending on the values in counts.
        '''
        outputs = []
        for i in range(num_rows):
            outputs.append(x[i, :counts[i]])
        return tf.concat(outputs, axis=0)
    
    # Positive anchors contribute to the loss, but negative and
    # neutral anchors (match value of 0 or -1) don't.
    indices = tf.where(tf.equal(target_matchs, 1))

    # Pick bbox deltas that contribute to the loss
    rpn_deltas = tf.gather_nd(rpn_deltas, indices)

    # Trim target bounding box deltas to the same length as rpn_deltas.
    batch_counts = tf.reduce_sum(tf.to_int32(tf.equal(target_matchs, 1)), axis=1)
    target_deltas = batch_pack(target_deltas, batch_counts,
                              target_deltas.shape.as_list()[0])

    loss = smooth_l1_loss(target_deltas, rpn_deltas)
    
    loss = tf.reduce_mean(loss) if tf.size(loss) > 0 else tf.constant(0.0)
    
    return loss
//  [python]
def rpn_bbox_loss(target_deltas, target_matchs, rpn_deltas):
    '''Return the RPN bounding box loss graph.
    
    Args
    ---
        target_deltas: [batch, num_rpn_deltas, (dy, dx, log(dh), log(dw))].
            Uses 0 padding to fill in unsed bbox deltas.
        target_matchs: [batch, anchors]. Anchor match type. 1=positive,
            -1=negative, 0=neutral anchor.
        rpn_deltas: [batch, anchors, (dy, dx, log(dh), log(dw))]
    '''
    def batch_pack(x, counts, num_rows):
        '''Picks different number of values from each row
        in x depending on the values in counts.
        '''
        outputs = []
        for i in range(num_rows):
            outputs.append(x[i, :counts[i]])
        return tf.concat(outputs, axis=0)
    
    # Positive anchors contribute to the loss, but negative and
    # neutral anchors (match value of 0 or -1) don't.
    indices = tf.where(tf.equal(target_matchs, 1))

    # Pick bbox deltas that contribute to the loss
    rpn_deltas = tf.gather_nd(rpn_deltas, indices)

    # Trim target bounding box deltas to the same length as rpn_deltas.
    batch_counts = tf.reduce_sum(tf.cast(tf.equal(target_matchs, 1), tf.int32), axis=1)
    target_deltas = batch_pack(target_deltas, batch_counts,
                              target_deltas.shape.as_list()[0])

    loss = smooth_l1_loss(target_deltas, rpn_deltas)
    
    loss = tf.reduce_mean(loss) if tf.size(loss) > 0 else tf.constant(0.0)
    
    return loss
//  [python]
def rcnn_class_loss(target_matchs_list, rcnn_class_logits_list):
    '''Loss for the classifier head of Faster RCNN.
    
    Args
    ---
        target_matchs_list: list of [num_rois]. Integer class IDs. Uses zero
            padding to fill in the array.
        rcnn_class_logits_list: list of [num_rois, num_classes]
    '''
    
    class_ids = tf.concat(target_matchs_list, 0)
    class_logits = tf.concat(rcnn_class_logits_list, 0)
    class_ids = tf.to_int64(class_ids)
    
    # loss = tf.losses.sparse_softmax_cross_entropy(labels=class_ids,
    #                                               logits=class_logits)

    num_classes = class_logits.shape[-1]
    # print(class_logits.shape)
    loss = keras.losses.categorical_crossentropy(tf.one_hot(class_ids, depth=num_classes),
                                                 class_logits, from_logits=True)


    loss = tf.reduce_mean(loss) if tf.size(loss) > 0 else tf.constant(0.0)
    return loss
//  [python]
def rcnn_class_loss(target_matchs_list, rcnn_class_logits_list):
    '''Loss for the classifier head of Faster RCNN.
    
    Args
    ---
        target_matchs_list: list of [num_rois]. Integer class IDs. Uses zero
            padding to fill in the array.
        rcnn_class_logits_list: list of [num_rois, num_classes]
    '''
    
    class_ids = tf.concat(target_matchs_list, 0)
    class_logits = tf.concat(rcnn_class_logits_list, 0)
    class_ids = tf.cast(class_ids, 'int64')
    
    # loss = tf.losses.sparse_softmax_cross_entropy(labels=class_ids,
    #                                               logits=class_logits)

    num_classes = class_logits.shape[-1]
    # print(class_logits.shape)
    loss = keras.losses.categorical_crossentropy(tf.one_hot(class_ids, depth=num_classes),
                                                 class_logits, from_logits=True)


    loss = tf.reduce_mean(loss) if tf.size(loss) > 0 else tf.constant(0.0)
    return loss
//  [python]
def inverse_sigmoid_decay(k, global_step_op):
  with tf.name_scope('inverse_sigmoid_decay'):
    k = tf.constant(k, dtype=tf.float32)
    tmp = k*tf.exp(-tf.to_float(global_step_op)/k)
    tmp = tmp / (1. + tmp)
  return tmp
//  [python]
def inverse_sigmoid_decay(k, global_step_op):
  with tf.name_scope('inverse_sigmoid_decay'):
    k = tf.constant(k, dtype=tf.float32)
    tmp = k*tf.exp(-tf.cast(global_step_op, tf.float32)/k)
    tmp = tmp / (1. + tmp)
  return tmp
//  [python]
def loss_function(self, pred, target, mask=None):
        """
        L2 loss function calculated with prediction and target.
        :param pred: The predicted image.
        :param target: The target image.
        :param mask: If not none, calculate loss only pixels, where mask == 1
        :return: L2 loss of (pred - target) / batch_size
        """
        batch_size, channel_size, image_size = get_batch_channel_image_size(pred, self.data_format, as_tensor=True)
        if mask is not None:
            diff = (pred - target) * mask
        else:
            diff = pred - target
        return tf.nn.l2_loss(diff) / tf.to_float(batch_size * 1024)
//  [python]
def loss_function(self, pred, target, mask=None):
        """
        L2 loss function calculated with prediction and target.
        :param pred: The predicted image.
        :param target: The target image.
        :param mask: If not none, calculate loss only pixels, where mask == 1
        :return: L2 loss of (pred - target) / batch_size
        """
        batch_size, channel_size, image_size = get_batch_channel_image_size(pred, self.data_format, as_tensor=True)
        if mask is not None:
            diff = (pred - target) * mask
        else:
            diff = pred - target
        return tf.nn.l2_loss(diff) / tf.cast(batch_size * 1024, tf.float32)
//  [python]
def _undiscretize_struct(struct, step_size, tf_dtype_struct):
  """Unscales the discretized structure and casts back to original dtypes."""

  def undiscretize_tensor(x, original_dtype):
    unscaled_x = t.cast(x, t.float32) * step_size
    return t.cast(unscaled_x, original_dtype)

  return tf.nest.map_structure(undiscretize_tensor, struct, tf_dtype_struct)
//  [python]
def _undiscretize_struct(struct, step_size, tf_dtype_struct):
  """Unscales the discretized structure and casts back to original dtypes."""

  def undiscretize_tensor(x, original_dtype):
    unscaled_x = tf.cast(x, tf.float32) * step_size
    return tf.cast(unscaled_x, original_dtype)

  return tf.nest.map_structure(undiscretize_tensor, struct, tf_dtype_struct)
//  [python]
def _fft_norm(
    x: Union[tf.Tensor, tf.Variable],
    dim: int,
    /,
    *,
    norm: str = "backward",
):
    n = tf.constant(x.shape[dim], dtype=x.dtype)
    if norm == "backward":
        return x
    elif norm == "ortho":
        return x / t.cast(t.sqrt(t.cast(n, tf.float32)), x.dtype)
    elif norm == "forward":
        return x / t.cast(n, x.dtype)
    else:
        raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}")
//  [python]
def _fft_norm(
    x: Union[tf.Tensor, tf.Variable],
    dim: int,
    /,
    *,
    norm: str = "backward",
):
    n = tf.constant(x.shape[dim], dtype=x.dtype)
    if norm == "backward":
        return x
    elif norm == "ortho":
        return x / tf.cast(tf.sqrt(tf.cast(n, tf.float32)), x.dtype)
    elif norm == "forward":
        return x / tf.cast(n, x.dtype)
    else:
        raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}")
//  [python]
def rcnn_bbox_loss(target_deltas_list, target_matchs_list, rcnn_deltas_list):
    '''Loss for Faster R-CNN bounding box refinement.
    
    Args
    ---
        target_deltas_list: list of [num_positive_rois, (dy, dx, log(dh), log(dw))]
        target_matchs_list: list of [num_rois]. Integer class IDs.
        rcnn_deltas_list: list of [num_rois, num_classes, (dy, dx, log(dh), log(dw))]
    '''
    
    target_deltas = tf.concat(target_deltas_list, 0)
    target_class_ids = tf.concat(target_matchs_list, 0)
    rcnn_deltas = tf.concat(rcnn_deltas_list, 0)

    # Only positive ROIs contribute to the loss. And only
    # the right class_id of each ROI. Get their indicies.
    positive_roi_ix = tf.where(target_class_ids > 0)[:, 0]
    positive_roi_class_ids = tf.to_int64(
        tf.gather(target_class_ids, positive_roi_ix))
    indices = tf.stack([positive_roi_ix, positive_roi_class_ids], axis=1)
    
    # Gather the deltas (predicted and true) that contribute to loss
    rcnn_deltas = tf.gather_nd(rcnn_deltas, indices)

    # Smooth-L1 Loss
    loss = smooth_l1_loss(target_deltas, rcnn_deltas)
    loss = tf.reduce_mean(loss) if tf.size(loss) > 0 else tf.constant(0.0)

    return loss
//  [python]
def rcnn_bbox_loss(target_deltas_list, target_matchs_list, rcnn_deltas_list):
    '''Loss for Faster R-CNN bounding box refinement.
    
    Args
    ---
        target_deltas_list: list of [num_positive_rois, (dy, dx, log(dh), log(dw))]
        target_matchs_list: list of [num_rois]. Integer class IDs.
        rcnn_deltas_list: list of [num_rois, num_classes, (dy, dx, log(dh), log(dw))]
    '''
    
    target_deltas = tf.concat(target_deltas_list, 0)
    target_class_ids = tf.concat(target_matchs_list, 0)
    rcnn_deltas = tf.concat(rcnn_deltas_list, 0)

    # Only positive ROIs contribute to the loss. And only
    # the right class_id of each ROI. Get their indicies.
    positive_roi_ix = tf.where(target_class_ids > 0)[:, 0]
    positive_roi_class_ids = tf.cast(
        tf.gather(target_class_ids, positive_roi_ix), tf.int64)
    indices = tf.stack([positive_roi_ix, positive_roi_class_ids], axis=1)
    
    # Gather the deltas (predicted and true) that contribute to loss
    rcnn_deltas = tf.gather_nd(rcnn_deltas, indices)

    # Smooth-L1 Loss
    loss = smooth_l1_loss(target_deltas, rcnn_deltas)
    loss = tf.reduce_mean(loss) if tf.size(loss) > 0 else tf.constant(0.0)

    return loss
//  [python]
def norm_initialization(norm, shape, x):
    if norm == "backward":
        norm_factor = tf.constant(1, x.dtype)
    elif norm in ["forward", "ortho"]:
        norm_factor = tf.cast(tf.math.reduce_prod(shape), x.dtype)
        if norm == "ortho":
            norm_factor = tf.math.sqrt(norm_factor)
    return norm_factor
//  [python]
def fft2(
    x: Union[tf.Tensor, tf.Variable],
    *,
    s: Optional[Sequence[int]] = None,
    dim: Sequence[int] = (-2, -1),
    norm: str = "backward",
    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
) -> Union[tf.Tensor, tf.Variable]:
    if s is None:
        s = (x.shape[dim[0]], x.shape[dim[1]])
    if len(x.shape) > 2:
        result = _fft2_helper(x, s, dim)
    else:
        x_new = trans_x_to_s(x, s, dim)
        x_complex = t.cast(x_new, tf.complex128)
        result = tf.signal.fft2d(x_complex)

    result = _fft2_norm(result, s, dim, norm)
    if x.dtype == tf.complex64:
        result = t.cast(result, dtype=tf.complex128)
    return result
//  [python]
def fft2(
    x: Union[tf.Tensor, tf.Variable],
    *,
    s: Optional[Sequence[int]] = None,
    dim: Sequence[int] = (-2, -1),
    norm: str = "backward",
    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
) -> Union[tf.Tensor, tf.Variable]:
    if s is None:
        s = (x.shape[dim[0]], x.shape[dim[1]])
    if len(x.shape) > 2:
        result = _fft2_helper(x, s, dim)
    else:
        x_new = trans_x_to_s(x, s, dim)
        x_complex = tf.cast(x_new, tf.complex128)
        result = tf.signal.fft2d(x_complex)

    result = _fft2_norm(result, s, dim, norm)
    if x.dtype == tf.complex64:
        result = tf.cast(result, dtype=tf.complex128)
    return result
//  [python]
def rfftn(
    x: Union[tf.Tensor, tf.Variable],
    s: Optional[Union[int, Tuple[int]]] = None,
    axes: Optional[Union[int, Tuple[int]]] = None,
    *,
    norm: str = "backward",
    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
) -> Union[tf.Tensor, tf.Variable]:
    result = _rfftn_helper(x, s, axes, norm)

    if out is not None:
        out = tf.cast(result, tf.complex128)
        # out = result
        return out
    else:
        # return result
        return tf.cast(result, tf.complex128)
//  [python]
def rfft_operations(x, rank, norm_factor):
    if x.shape.rank == 1:
        x = tf.signal.rfft(x)
    elif x.shape.rank == 2:
        x = tf.switch_case(
            rank - 1, {0: lambda: tf.signal.rfft(x), 1: lambda: tf.signal.rfft2d(x)}
        )
    else:
        x = tf.switch_case(
            rank - 1,
            {
                0: lambda: tf.signal.rfft(x),
                1: lambda: tf.signal.rfft2d(x),
                2: lambda: tf.signal.rfft3d(x),
            },
        )
    norm_factor = t.cast(norm_factor, tf.complex128)
    x = t.cast(x, tf.complex128)
    x = x / norm_factor
    return x
//  [python]
def rfft_operations(x, rank, norm_factor):
    if x.shape.rank == 1:
        x = tf.signal.rfft(x)
    elif x.shape.rank == 2:
        x = tf.switch_case(
            rank - 1, {0: lambda: tf.signal.rfft(x), 1: lambda: tf.signal.rfft2d(x)}
        )
    else:
        x = tf.switch_case(
            rank - 1,
            {
                0: lambda: tf.signal.rfft(x),
                1: lambda: tf.signal.rfft2d(x),
                2: lambda: tf.signal.rfft3d(x),
            },
        )
    norm_factor = tf.cast(norm_factor, tf.complex128)
    x = tf.cast(x, tf.complex128)
    x = x / norm_factor
    return x
//  [python]
def fft(
    x: Union[tf.Tensor, tf.Variable],
    dim: int,
    /,
    *,
    norm: str = "backward",
    n: Optional[Union[int, Tuple[int]]] = None,
    out: Optional[Union[tf.Tensor, tf.Variable]] = None,
) -> Union[tf.Tensor, tf.Variable]:
    # ToDo: Remove conversion from float to complex when casting mode is working
    if x.dtype == "float32":
        x = tf.cast(x, tf.complex64)
    elif x.dtype == "float64":
        x = tf.cast(x, tf.complex128)
    if not isinstance(dim, int):
        raise ivy.utils.exceptions.IvyError(
            f"Expecting <class 'int'> instead of {type(dim)}"
        )
    if n is None:
        n = x.shape[dim]
    if n < -len(x.shape):
        raise ivy.utils.exceptions.IvyError(
            f"Invalid dim {dim}, expecting ranging"
            " from {-len(x.shape)} to {len(x.shape)-1}  "
        )
    if not isinstance(n, int):
        raise ivy.utils.exceptions.IvyError(
            f"Expecting <class 'int'> instead of {type(n)}"
        )
    if n <= 1:
        raise ivy.utils.exceptions.IvyError(
            f"Invalid data points {n}, expecting more than 1"
        )
    if norm not in ["backward", "ortho", "forward"]:
        raise ivy.utils.exceptions.IvyError(f"Unrecognized normalization mode {norm}")
    if x.shape[dim] != n:
        s = list(x.shape)
        if s[dim] > n:
            index = [slice(None)] * len(s)
            index[dim] = slice(0, n)
            x = x[tuple(index)]
            del index
        else:
            s[dim] = n - s[dim]
            z = tf.zeros(s, x.dtype)
            x = tf.concat([x, z], dim)
        del s
    operation_name = f"{n} points FFT at dim {dim} with {norm} normalization"
    if dim != -1 or dim != len(x.shape) - 1:
        permute = [i for i in range(len(x.shape))]
        permute[dim], permute[-1] = permute[-1], permute[dim]
        x = tf.transpose(x, permute)
        ret = tf.signal.fft(x, operation_name)
        ret = tf.transpose(ret, permute)
        del permute
    else:
        ret = tf.signal.fft(x, operation_name)
    ret = _fft_norm(ret, dim, norm=norm)
    return ret
//  [python]
def compute_v_x(self, v_obs, x_obs):
        """Equation (33)"""

        # Compute probability mass function for the symbols
        x_obs = tf.expand_dims(x_obs, -1)
        v_obs = tf.expand_dims(v_obs, -1)

        points = expand_to_rank(self._points, tf.rank(x_obs), axis=0)
        logits = -tf.pow(x_obs-points, 2) / (tf.to_float(2)*v_obs)
        pmf = tf.math.softmax(logits)

        # Compute mean and variance of all symbols
        x = tf.reduce_sum(points * pmf, axis=-1, keepdims=True)
        v = tf.reduce_sum((points-x)**2 * pmf, axis=-1)
        v = tf.maximum(v, self._prec)
        x = tf.squeeze(x, axis=-1)

        return v, x, logits
//  [python]
def compute_v_x(self, v_obs, x_obs):
        """Equation (33)"""

        # Compute probability mass function for the symbols
        x_obs = tf.expand_dims(x_obs, -1)
        v_obs = tf.expand_dims(v_obs, -1)

        points = expand_to_rank(self._points, tf.rank(x_obs), axis=0)
        logits = -tf.pow(x_obs-points, 2) / (tf.cast(2, self._rdtype)*v_obs)
        pmf = tf.math.softmax(logits)

        # Compute mean and variance of all symbols
        x = tf.reduce_sum(points * pmf, axis=-1, keepdims=True)
        v = tf.reduce_sum((points-x)**2 * pmf, axis=-1)
        v = tf.maximum(v, self._prec)
        x = tf.squeeze(x, axis=-1)

        return v, x, logits
//  [python]
def trans_x_to_s(
    x: Union[tf.Tensor, tf.Variable],
    s: Optional[Sequence[int]] = None,
    dim: Sequence[int] = (-2, -1),
) -> Union[tf.Tensor, tf.Variable]:
    """Change the shape of the input array x to the desired output shape s."""
    if x.dtype not in [tf.complex128, tf.complex64]:
        x = t.cast(x, tf.float32)
    x_shape = x.shape
    if dim in [(-1, -2), (1, 0)]:
        s = (s[1], s[0])
    if s[0] >= x_shape[0] and s[1] >= x_shape[1]:
        paddings = tf.constant([[0, s[0] - x_shape[0]], [0, s[1] - x_shape[1]]])
        x_new = tf.pad(x, paddings=paddings)
    elif (s[0] <= x_shape[0] or s[1] <= x_shape[1]) and min(s) > min(x_shape):
        x_new = x[: s[0], : s[1]]
        if s[0] != x_new.shape[0]:
            size = s[0] - x_new.shape[0]
            z = tf.zeros((size, s[1]), dtype=x.dtype)
            x_new = tf.concat([x_new, z], 0)
        elif s[1] != x_new.shape[1]:
            size = s[1] - x_new.shape[1]
            z = tf.zeros((s[0], size), dtype=x.dtype)
            x_new = tf.concat([x_new, z], 1)
    elif (s[0] >= x_shape[0] and s[1] <= x_shape[1]) and min(s) <= min(x_shape):
        x_new = x[: s[0], : s[1]]
        size = s[0] - x_new.shape[0]
        z = tf.zeros((size, s[1]), dtype=x.dtype)
        x_new = tf.concat([x_new, z], 0)
    elif (s[0] < x_shape[0] and s[1] > x_shape[1]) and min(s) == min(x_shape):
        x_new = x[: s[0], : s[1]]
        size = s[1] - x_new.shape[1]
        z = tf.zeros((s[0], size), dtype=x.dtype)
        x_new = tf.concat([x_new, z], axis=1)
    else:
        x_new = x[: s[0], : s[1]]
    return x_new
//  [python]
def trans_x_to_s(
    x: Union[tf.Tensor, tf.Variable],
    s: Optional[Sequence[int]] = None,
    dim: Sequence[int] = (-2, -1),
) -> Union[tf.Tensor, tf.Variable]:
    """Change the shape of the input array x to the desired output shape s."""
    if x.dtype not in [tf.complex128, tf.complex64]:
        x = tf.cast(x, tf.float32)
    x_shape = x.shape
    if dim in [(-1, -2), (1, 0)]:
        s = (s[1], s[0])
    if s[0] >= x_shape[0] and s[1] >= x_shape[1]:
        paddings = tf.constant([[0, s[0] - x_shape[0]], [0, s[1] - x_shape[1]]])
        x_new = tf.pad(x, paddings=paddings)
    elif (s[0] <= x_shape[0] or s[1] <= x_shape[1]) and min(s) > min(x_shape):
        x_new = x[: s[0], : s[1]]
        if s[0] != x_new.shape[0]:
            size = s[0] - x_new.shape[0]
            z = tf.zeros((size, s[1]), dtype=x.dtype)
            x_new = tf.concat([x_new, z], 0)
        elif s[1] != x_new.shape[1]:
            size = s[1] - x_new.shape[1]
            z = tf.zeros((s[0], size), dtype=x.dtype)
            x_new = tf.concat([x_new, z], 1)
    elif (s[0] >= x_shape[0] and s[1] <= x_shape[1]) and min(s) <= min(x_shape):
        x_new = x[: s[0], : s[1]]
        size = s[0] - x_new.shape[0]
        z = tf.zeros((size, s[1]), dtype=x.dtype)
        x_new = tf.concat([x_new, z], 0)
    elif (s[0] < x_shape[0] and s[1] > x_shape[1]) and min(s) == min(x_shape):
        x_new = x[: s[0], : s[1]]
        size = s[1] - x_new.shape[1]
        z = tf.zeros((s[0], size), dtype=x.dtype)
        x_new = tf.concat([x_new, z], axis=1)
    else:
        x_new = x[: s[0], : s[1]]
    return x_new