jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def iou_coe(output, target, threshold=0.5, axis=[1,2,3], smooth=1e-5):
    """Non-differentiable Intersection over Union (IoU) for comparing the
    similarity of two batch of data, usually be used for evaluating binary image segmentation.
    The coefficient between 0 to 1, 1 means totally match.

    Parameters
    -----------
    output : tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    target : tensor
        A distribution with shape: [batch_size, ....], (any dimensions).
    threshold : float
        The threshold value to be true.
    axis : list of integer
        All dimensions are reduced, default ``[1,2,3]``.
    smooth : float
        This small value will be added to the numerator and denominator, see ``dice_coe``.

    Notes
    ------
    - IoU cannot be used as training loss, people usually use dice coefficient for training, IoU and hard-dice for evaluating.
    """
    pre = tf.cast(output > threshold, dtype=tf.float32)
    truth = tf.cast(target > threshold, dtype=tf.float32)
    inse = tf.reduce_sum(tf.multiply(pre, truth), axis=axis) # AND
    union = tf.reduce_sum(tf.cast(tf.add(pre, truth)>= 1, dtype=tf.float32), axis=axis) # OR
    ## old axis=[0,1,2,3]
    # epsilon = 1e-5
    # batch_iou = inse / (union + epsilon)
    ## new haodong
    batch_iou = (inse + smooth) / (union + smooth)
    iou = tf.reduce_mean(batch_iou)
    return iou
//  [python]
def modulated_conv2d_layer(x, y, fmaps, kernel, up=False, down=False, demodulate=True, resample_kernel=None, gain=1, use_wscale=True, lrmul=1, fused_modconv=True, weight_var='weight', mod_weight_var='mod_weight', mod_bias_var='mod_bias', adapt_func='training.networks_stylegan2.apply_identity', rho_in=1, svd=False, factorized=False, sv_factors=None, svd_center=False, lambda_mask=None, reduce_dims=None, spatial=False, svd_config='S'):
    assert not (up and down)
    assert kernel >= 1 and kernel % 2 == 1
    x_prev = x
    w = get_svd_weight([kernel, kernel, x.shape[1].value, fmaps], 
                       gain=gain, use_wscale=use_wscale, lrmul=lrmul, weight_var=weight_var, 
                       spatial=spatial, svd=svd, factorized=factorized, sv_factors=sv_factors, svd_center=svd_center, lambda_mask=lambda_mask, reduce_dims=reduce_dims,
                       svd_config=svd_config)
    # w = _get_conv_w(x, fmaps, kernel, gain, use_wscale, lrmul, weight_var, svd, factorized, sv_factors, lambda_mask)
    ww = w[np.newaxis] # [BkkIO] Introduce minibatch dimension.

    # Modulate.
    s = dense_layer(y, fmaps=x.shape[1].value, weight_var=mod_weight_var) # [BI] Transform incoming W to style.
    s = apply_bias_act(s, bias_var=mod_bias_var) + 1 # [BI] Add bias (initially 1).
    ww *= t.cast(s[:, np.newaxis, np.newaxis, :, np.newaxis], w.dtype) # [BkkIO] Scale input feature maps.

    # Demodulate.
    if demodulate:
        d = tf.rsqrt(tf.reduce_sum(tf.square(ww), axis=[1,2,3]) + 1e-8) # [BO] Scaling factor.
        ww *= d[:, np.newaxis, np.newaxis, np.newaxis, :] # [BkkIO] Scale output feature maps.

    # Reshape/scale input.
    if fused_modconv:
        w = tf.reshape(tf.transpose(ww, [1, 2, 3, 0, 4]), [ww.shape[1], ww.shape[2], ww.shape[3], -1])
        x = tf.reshape(x, [1, -1, x.shape[2], x.shape[3]]) # Fused => reshape minibatch to convolution groups.
    else:
        x *= t.cast(s[:, :, np.newaxis, np.newaxis], x.dtype) # [BIhw] Not fused => scale input activations.

    # Convolution with optional up/downsampling.
    if up:
        x = upsample_conv_2d(x, t.cast(w, x.dtype), data_format='NCHW', k=resample_kernel)
    elif down:
        x = conv_downsample_2d(x, t.cast(w, x.dtype), data_format='NCHW', k=resample_kernel)
    else:
        x = tf.nn.conv2d(x, t.cast(w, x.dtype), data_format='NCHW', strides=[1,1,1,1], padding='SAME')

    # Reshape/scale output.
    if fused_modconv:
        x = tf.reshape(x, [-1, fmaps, x.shape[2], x.shape[3]]) # Fused => reshape convolution groups back to minibatch.
    elif demodulate:
        x *= t.cast(d[:, :, np.newaxis, np.newaxis], x.dtype) # [BOhw] Not fused => scale output activations.

    x = call_func_by_name(func_name=adapt_func, x=x, x_prev=x_prev, rho_in=rho_in)
    return x
//  [python]
def modulated_conv2d_layer(x, y, fmaps, kernel, up=False, down=False, demodulate=True, resample_kernel=None, gain=1, use_wscale=True, lrmul=1, fused_modconv=True, weight_var='weight', mod_weight_var='mod_weight', mod_bias_var='mod_bias', adapt_func='training.networks_stylegan2.apply_identity', rho_in=1, svd=False, factorized=False, sv_factors=None, svd_center=False, lambda_mask=None, reduce_dims=None, spatial=False, svd_config='S'):
    assert not (up and down)
    assert kernel >= 1 and kernel % 2 == 1
    x_prev = x
    w = get_svd_weight([kernel, kernel, x.shape[1].value, fmaps], 
                       gain=gain, use_wscale=use_wscale, lrmul=lrmul, weight_var=weight_var, 
                       spatial=spatial, svd=svd, factorized=factorized, sv_factors=sv_factors, svd_center=svd_center, lambda_mask=lambda_mask, reduce_dims=reduce_dims,
                       svd_config=svd_config)
    # w = _get_conv_w(x, fmaps, kernel, gain, use_wscale, lrmul, weight_var, svd, factorized, sv_factors, lambda_mask)
    ww = w[np.newaxis] # [BkkIO] Introduce minibatch dimension.

    # Modulate.
    s = dense_layer(y, fmaps=x.shape[1].value, weight_var=mod_weight_var) # [BI] Transform incoming W to style.
    s = apply_bias_act(s, bias_var=mod_bias_var) + 1 # [BI] Add bias (initially 1).
    ww *= tf.cast(s[:, np.newaxis, np.newaxis, :, np.newaxis], w.dtype) # [BkkIO] Scale input feature maps.

    # Demodulate.
    if demodulate:
        d = tf.rsqrt(tf.reduce_sum(tf.square(ww), axis=[1,2,3]) + 1e-8) # [BO] Scaling factor.
        ww *= d[:, np.newaxis, np.newaxis, np.newaxis, :] # [BkkIO] Scale output feature maps.

    # Reshape/scale input.
    if fused_modconv:
        w = tf.reshape(tf.transpose(ww, [1, 2, 3, 0, 4]), [ww.shape[1], ww.shape[2], ww.shape[3], -1])
        x = tf.reshape(x, [1, -1, x.shape[2], x.shape[3]]) # Fused => reshape minibatch to convolution groups.
    else:
        x *= tf.cast(s[:, :, np.newaxis, np.newaxis], x.dtype) # [BIhw] Not fused => scale input activations.

    # Convolution with optional up/downsampling.
    if up:
        x = upsample_conv_2d(x, tf.cast(w, x.dtype), data_format='NCHW', k=resample_kernel)
    elif down:
        x = conv_downsample_2d(x, tf.cast(w, x.dtype), data_format='NCHW', k=resample_kernel)
    else:
        x = tf.nn.conv2d(x, tf.cast(w, x.dtype), data_format='NCHW', strides=[1,1,1,1], padding='SAME')

    # Reshape/scale output.
    if fused_modconv:
        x = tf.reshape(x, [-1, fmaps, x.shape[2], x.shape[3]]) # Fused => reshape convolution groups back to minibatch.
    elif demodulate:
        x *= tf.cast(d[:, :, np.newaxis, np.newaxis], x.dtype) # [BOhw] Not fused => scale output activations.

    x = call_func_by_name(func_name=adapt_func, x=x, x_prev=x_prev, rho_in=rho_in)
    return x
//  [python]
def __call__(self, shape, dtype=None):
        if dtype is None:
            dtype = tf.float32

        with tf.device('/CPU:0'):
            maxval = tf.sqrt(tf.constant(1.) / tf.cast(shape[0], tf.float32))
            maxval = tf.to_float(maxval)
            minval = -maxval

            weights = tf.random.uniform(shape, minval=minval, maxval=maxval, dtype=dtype)
            weights = tf.to_float(weights)
        return weights
//  [python]
def __call__(self, shape, dtype=None):
        if dtype is None:
            dtype = tf.float32

        with tf.device('/CPU:0'):
            maxval = tf.sqrt(tf.constant(1.) / tf.cast(shape[0], tf.float32))
            maxval = tf.cast(maxval, dtype=dtype)
            minval = -maxval

            weights = tf.random.uniform(shape, minval=minval, maxval=maxval, dtype=dtype)
            weights = tf.cast(weights, dtype=dtype)
        return weights
//  [python]
def _shard_inputs(self, features, num_towers):
    """docstring."""
    if num_towers == 1:
      return [features]

    source = features["source"]
    target_input = features["target_input"]
    target_output = features["target_output"]
    source_seq_length = features["source_sequence_length"]
    target_seq_length = features["target_sequence_length"]

    # Compute each split sizes.
    global_batch_size = tf.size(source_seq_length)
    tower_batch_size = tf.to_int32(global_batch_size / num_towers)

    split_sizes = [tower_batch_size] * (num_towers - 1)
    split_sizes.append(global_batch_size - (num_towers - 1) * tower_batch_size)

    sources = tf.split(source, split_sizes, axis=0)
    target_inputs = tf.split(target_input, split_sizes, axis=0)
    target_outputs = tf.split(target_output, split_sizes, axis=0)
    source_sequence_lengths = tf.split(source_seq_length, split_sizes)
    target_sequence_lengths = tf.split(target_seq_length, split_sizes)

    tower_features = []
    for i in range(num_towers):
      tower_features.append({
          "source": sources[i],
          "target_input": target_inputs[i],
          "target_output": target_outputs[i],
          "source_sequence_length": source_sequence_lengths[i],
          "target_sequence_length": target_sequence_lengths[i]
      })

    return tower_features
//  [python]
def _shard_inputs(self, features, num_towers):
    """docstring."""
    if num_towers == 1:
      return [features]

    source = features["source"]
    target_input = features["target_input"]
    target_output = features["target_output"]
    source_seq_length = features["source_sequence_length"]
    target_seq_length = features["target_sequence_length"]

    # Compute each split sizes.
    global_batch_size = tf.size(source_seq_length)
    tower_batch_size = tf.cast(global_batch_size / num_towers, dtype=tf.int32)

    split_sizes = [tower_batch_size] * (num_towers - 1)
    split_sizes.append(global_batch_size - (num_towers - 1) * tower_batch_size)

    sources = tf.split(source, split_sizes, axis=0)
    target_inputs = tf.split(target_input, split_sizes, axis=0)
    target_outputs = tf.split(target_output, split_sizes, axis=0)
    source_sequence_lengths = tf.split(source_seq_length, split_sizes)
    target_sequence_lengths = tf.split(target_seq_length, split_sizes)

    tower_features = []
    for i in range(num_towers):
      tower_features.append({
          "source": sources[i],
          "target_input": target_inputs[i],
          "target_output": target_outputs[i],
          "source_sequence_length": source_sequence_lengths[i],
          "target_sequence_length": target_sequence_lengths[i]
      })

    return tower_features
//  [python]
def _get_custom_getter():
  """Returns a custom getter that this class's methods must be called under.

  All methods of this class must be called under a variable scope that was
  passed this custom getter. Example:

  
  network = ConvNetBuilder(...)
  with tf.variable_scope('cg', custom_getter=network.get_custom_getter()):
    network.conv(...)
    # Call more methods of network here
  

  Currently, this custom getter only does anything if self.use_tf_layers is
  True. In that case, it causes variables to be stored as dtype
  self.variable_type, then casted to the requested dtype, instead of directly
  storing the variable as the requested dtype.
  """

  def inner_custom_getter(getter, *args, **kwargs):
    """Custom getter that forces variables to have type self.variable_type."""
    cast_to_float16 = False
    requested_dtype = kwargs["dtype"]
    if requested_dtype == tf.float16:
      # Only change the variable dtype if doing so does not decrease variable
      # precision.
      kwargs["dtype"] = tf.float32
      cast_to_float16 = True
    var = getter(*args, **kwargs)
    with tf_ops.init_scope():
      # This if statement is needed to guard the cast, because batch norm
      # assigns directly to the return value of this custom getter. The cast
      # makes the return value not a variable so it cannot be assigned. Batch
      # norm variables are always in fp32 so this if statement is never
      # triggered for them.
      if cast_to_float16:
        var = tf.cast(var, tf.float16)
    return var

  return inner_custom_getter
//  [python]
def _get_custom_getter():
  """Returns a custom getter that this class's methods must be called under.

  All methods of this class must be called under a variable scope that was
  passed this custom getter. Example:

  ```python
  network = ConvNetBuilder(...)
  with tf.variable_scope('cg', custom_getter=network.get_custom_getter()):
    network.conv(...)
    # Call more methods of network here
  ```

  Currently, this custom getter only does anything if self.use_tf_layers is
  True. In that case, it causes variables to be stored as dtype
  self.variable_type, then casted to the requested dtype, instead of directly
  storing the variable as the requested dtype.
  """

  def inner_custom_getter(getter, *args, **kwargs):
    """Custom getter that forces variables to have type self.variable_type."""
    cast_to_float16 = False
    requested_dtype = kwargs["dtype"]
    if requested_dtype == tf.float16:
      # Only change the variable dtype if doing so does not decrease variable
      # precision.
      kwargs["dtype"] = tf.float32
      cast_to_float16 = True
    var = getter(*args, **kwargs)
    with tf_ops.init_scope():
      # This if statement is needed to guard the cast, because batch norm
      # assigns directly to the return value of this custom getter. The cast
      # makes the return value not a variable so it cannot be assigned. Batch
      # norm variables are always in fp32 so this if statement is never
      # triggered for them.
      if cast_to_float16:
        var = tf.cast(var, tf.float16)
    return var

  return inner_custom_getter
//  [python]
def call(self, indices):
        cpu_indices, gpu_indices = [], []

        if not self.cpu_inputs:
            return self.gpu_embedding(indices)

        if not self.gpu_inputs:
            with tf.device('/CPU:0'):
                return self.cpu_embedding(indices)

        for i in self.cpu_inputs:
            cpu_indices.append(indices[:, i])
        for i in self.gpu_inputs:
            gpu_indices.append(indices[:, i])

        to_concat = []
        # at least one cpu-based embedding
        with tf.device('/CPU:0'):
            cpu_indices = tf.stack(cpu_indices, axis=1)
            cpu_results = self.cpu_embedding(cpu_indices)
            cpu_results = tf.to_float(cpu_results)
            to_concat.append(cpu_results)
        # at least one gpu-based embedding
        with tf.device('/GPU:0'):
            gpu_indices = tf.stack(gpu_indices, axis=1)
            gpu_results = self.gpu_embedding(gpu_indices)
            to_concat.append(gpu_results)

        result = tf.concat(to_concat, axis=1)

        reorder_indices = np.concatenate([self.cpu_inputs, self.gpu_inputs], axis=0).argsort().tolist()
        split_result = tf.split(result, num_or_size_splits=indices.shape[1], axis=1)
        result = [split_result[i] for i in reorder_indices]
        result = tf.concat(result, axis=1)
        return result
//  [python]
def call(self, indices):
        cpu_indices, gpu_indices = [], []

        if not self.cpu_inputs:
            return self.gpu_embedding(indices)

        if not self.gpu_inputs:
            with tf.device('/CPU:0'):
                return self.cpu_embedding(indices)

        for i in self.cpu_inputs:
            cpu_indices.append(indices[:, i])
        for i in self.gpu_inputs:
            gpu_indices.append(indices[:, i])

        to_concat = []
        # at least one cpu-based embedding
        with tf.device('/CPU:0'):
            cpu_indices = tf.stack(cpu_indices, axis=1)
            cpu_results = self.cpu_embedding(cpu_indices)
            cpu_results = tf.cast(cpu_results, dtype=self.dtype)
            to_concat.append(cpu_results)
        # at least one gpu-based embedding
        with tf.device('/GPU:0'):
            gpu_indices = tf.stack(gpu_indices, axis=1)
            gpu_results = self.gpu_embedding(gpu_indices)
            to_concat.append(gpu_results)

        result = tf.concat(to_concat, axis=1)

        reorder_indices = np.concatenate([self.cpu_inputs, self.gpu_inputs], axis=0).argsort().tolist()
        split_result = tf.split(result, num_or_size_splits=indices.shape[1], axis=1)
        result = [split_result[i] for i in reorder_indices]
        result = tf.concat(result, axis=1)
        return result
//  [python]
def class_balanced_sigmoid_cross_entropy(logits, label, name='cross_entropy_loss'):
    """
    The class-balanced cross entropy loss,
    as in `Holistically-Nested Edge Detection
    <http://arxiv.org/abs/1504.06375>`_.

    Args:
        logits: of shape (b, ...).
        label: of the same shape. the ground truth in {0,1}.
    Returns:
        class-balanced cross entropy loss.
    """
    with tf.name_scope('class_balanced_sigmoid_cross_entropy'):
        y = tf.cast(label, tf.float32)

        count_neg = tf.reduce_sum(1. - y)
        count_pos = tf.reduce_sum(y)
        beta = count_neg / (count_neg + count_pos)

        pos_weight = beta / (1 - beta)
        cost = tf.nn.weighted_cross_entropy_with_logits(logits=logits, targets=y, pos_weight=pos_weight)
        cost = tf.reduce_mean(cost * (1 - beta))
        zero = tf.equal(count_pos, 0.0)
    return tf.where(zero, 0.0, cost, name=name)
//  [python]
def parse_csv_function(line):
  """
  Decode CSV for Dataset.
  
  Args:
    line: One line data of the CSV.
  
  Return:
    The op of features and labels
  """

  FIELD_DEFAULTS = [[0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0],
                    [0.0], [0]]

  fields = tf.decode_csv(line, FIELD_DEFAULTS)

  label = fields[-1]
  label = tf.to_int64(label)
  features = tf.stack(fields[0:-1])

  return features, label
//  [python]
def parse_csv_function(line):
  """
  Decode CSV for Dataset.
  
  Args:
    line: One line data of the CSV.
  
  Return:
    The op of features and labels
  """

  FIELD_DEFAULTS = [[0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0], [0.0],
                    [0.0], [0]]

  fields = tf.decode_csv(line, FIELD_DEFAULTS)

  label = fields[-1]
  label = tf.cast(label, tf.int64)
  features = tf.stack(fields[0:-1])

  return features, label
//  [python]
def update_state(self,
                   y_true,
                   y_pred,
                   sample_weight):
    batch = tf.shape(y_pred)[0]
    num_embs = tf.shape(y_pred)[1]
    proba_pred = tf.nn.sigmoid(y_pred)
    if sample_weight is not None:
      proba_pred *= sample_weight[Ellipsis, None]

    weights_range = tf.range(num_embs, dtype=tf.float32)
    weights_range_square = tf.abs(weights_range[:, tf.newaxis] -
                                  weights_range[tf.newaxis, :])
    indic_range_fun = lambda x, a, b: tf.logical_and(x >= a, x <= b)
    weights_square = indic_range_fun(weights_range_square,
                                     self._range_low,
                                     self._range_high)
    weights_square = tf.cast(weights_square[None, Ellipsis, None], dtype=tf.float32)

    proba_pred_filter = proba_pred * weights_square
    flat_proba_pred_filter = tf.reshape(proba_pred_filter, (batch, -1))
    y_true_filter = y_true * weights_square
    flat_y_true_filter = tf.reshape(y_true_filter, (batch, -1))

    _, indices = tf.math.top_k(flat_proba_pred_filter, k=self._at_k)
    flat_y_pred_filter = tf.to_float(flat_proba_pred_filter > 0.5)

    true_in_top = tf.gather(flat_y_true_filter, indices, batch_dims=-1)
    pred_in_top = tf.gather(flat_y_pred_filter, indices, batch_dims=-1)
    true_pred_in_top = tf.gather(
        flat_y_true_filter * flat_y_pred_filter, indices, batch_dims=-1)

    number_true = tf.maximum(tf.reduce_sum(true_in_top, axis=-1), 1e-6)
    number_preds = tf.maximum(tf.reduce_sum(pred_in_top, axis=-1), 1e-6)
    number_true_preds = tf.reduce_sum(true_pred_in_top, axis=-1)

    precision = tf.maximum(number_true_preds / number_preds, 1e-6)
    recall = tf.maximum(number_true_preds / number_true, 1e-6)

    self._precision.update_state(precision)
    self._recall.update_state(recall)
    self._f1score.update_state(
        2 * (precision * recall) / (precision + recall))

    # TODO(qberthet): double-check.
    auprc_weight = sample_weight * weights_square[Ellipsis, 0]
    self._auprc.update_state(
        y_true_filter, proba_pred_filter, sample_weight=auprc_weight)
//  [python]
def update_state(self,
                   y_true,
                   y_pred,
                   sample_weight):
    batch = tf.shape(y_pred)[0]
    num_embs = tf.shape(y_pred)[1]
    proba_pred = tf.nn.sigmoid(y_pred)
    if sample_weight is not None:
      proba_pred *= sample_weight[Ellipsis, None]

    weights_range = tf.range(num_embs, dtype=tf.float32)
    weights_range_square = tf.abs(weights_range[:, tf.newaxis] -
                                  weights_range[tf.newaxis, :])
    indic_range_fun = lambda x, a, b: tf.logical_and(x >= a, x <= b)
    weights_square = indic_range_fun(weights_range_square,
                                     self._range_low,
                                     self._range_high)
    weights_square = tf.cast(weights_square[None, Ellipsis, None], dtype=tf.float32)

    proba_pred_filter = proba_pred * weights_square
    flat_proba_pred_filter = tf.reshape(proba_pred_filter, (batch, -1))
    y_true_filter = y_true * weights_square
    flat_y_true_filter = tf.reshape(y_true_filter, (batch, -1))

    _, indices = tf.math.top_k(flat_proba_pred_filter, k=self._at_k)
    flat_y_pred_filter = tf.cast(flat_proba_pred_filter > 0.5, tf.float32)

    true_in_top = tf.gather(flat_y_true_filter, indices, batch_dims=-1)
    pred_in_top = tf.gather(flat_y_pred_filter, indices, batch_dims=-1)
    true_pred_in_top = tf.gather(
        flat_y_true_filter * flat_y_pred_filter, indices, batch_dims=-1)

    number_true = tf.maximum(tf.reduce_sum(true_in_top, axis=-1), 1e-6)
    number_preds = tf.maximum(tf.reduce_sum(pred_in_top, axis=-1), 1e-6)
    number_true_preds = tf.reduce_sum(true_pred_in_top, axis=-1)

    precision = tf.maximum(number_true_preds / number_preds, 1e-6)
    recall = tf.maximum(number_true_preds / number_true, 1e-6)

    self._precision.update_state(precision)
    self._recall.update_state(recall)
    self._f1score.update_state(
        2 * (precision * recall) / (precision + recall))

    # TODO(qberthet): double-check.
    auprc_weight = sample_weight * weights_square[Ellipsis, 0]
    self._auprc.update_state(
        y_true_filter, proba_pred_filter, sample_weight=auprc_weight)
//  [python]
def __declare_placeholders(self):
        # Placeholders for input and ground truth output
        self.input_batch = tf.placeholder(dtype=tf.int32, shape=[None, None], name='input_batch')
        self.ground_truth_tags = tf.placeholder(dtype=tf.int32, shape=[None, None], name='ground_truth_tags')

        # Placeholder for lengths of the sequences
        self.lengths = tf.placeholder(dtype=tf.int32, shape=[None], name='lengths')

        # Placeholder for a dropout keep probability. If we don't feed
        # a value for this placeholder, it will be equal to 1.0.
        self.dropout_ph = tf.placeholder_with_default(tf.constant(1.0, dtype=tf.float32), shape=[])

        # Placeholder for a learning rate (tf.float32)
        self.learning_rate_ph = tf.placeholder(dtype=tf.float32, shape=[])
//  [python]
def __declare_placeholders(self):
        # Placeholders for input and ground truth output
        self.input_batch = tf.placeholder(dtype=tf.int32, shape=[None, None], name='input_batch')
        self.ground_truth_tags = tf.placeholder(dtype=tf.int32, shape=[None, None], name='ground_truth_tags')

        # Placeholder for lengths of the sequences
        self.lengths = tf.placeholder(dtype=tf.int32, shape=[None], name='lengths')

        # Placeholder for a dropout keep probability. If we don't feed
        # a value for this placeholder, it will be equal to 1.0.
        self.dropout_ph = tf.placeholder_with_default(tf.cast(1.0, tf.float32), shape=[])

        # Placeholder for a learning rate (tf.float32)
        self.learning_rate_ph = tf.placeholder(dtype=tf.float32, shape=[])
//  [python]
def __optimize(self):
        """
        Specifies the optimizer and train_op for the model.

        :return:
        """
        self.optimizer = tf.train.AdamOptimizer(self.learning_rate_ph)
        self.grads_and_vars = self.optimizer.compute_gradients(self.loss)

        # Gradient clipping for self.grads_and_vars to eliminate exploding gradients.
        # Note that you need to apply this operation only for gradients,
        # self.grads_and_vars also contains variables
        clip_norm = tf.cast(1.0, tf.float32)
        self.grads_and_vars = [(tf.clip_by_norm(grad, clip_norm), var) for grad, var in self.grads_and_vars]
        self.train_op = self.optimizer.apply_gradients(self.grads_and_vars)
//  [python]
def call(self, inputs):
    if isinstance(inputs, (list, tuple)) and len(inputs) == 2:
      input_tensor, attention_mask = inputs
    else:
      input_tensor, attention_mask = (inputs, None)

    if self._output_range:
      target_tensor = input_tensor[:, 0:self._output_range, :]
      attention_mask = attention_mask[:, 0:self._output_range, :]
    else:
      if self._norm_first:
        source_tensor = input_tensor
        input_tensor = self._attention_layer_norm(input_tensor)
      target_tensor = input_tensor

    attention_output = self._attention_layer(
        query=target_tensor, value=input_tensor, attention_mask=attention_mask)
    attention_output = self._attention_dropout(attention_output)
    if self._norm_first:
      attention_output = source_tensor + attention_output
    else:
      attention_output = self._attention_layer_norm(target_tensor +
                                                    attention_output)
    if self._norm_first:
      source_attention_output = attention_output
      attention_output = self._output_layer_norm(attention_output)

    layer_output = self._output_dense(attention_output)
    layer_output = self._output_dropout(layer_output)
    # During mixed precision training, attention_output is from layer norm and
    # is always fp32 for now. Cast layer_output to fp32 for the subsequent
    # add.
    layer_output = t.cast(layer_output, t.float32)
    if self._norm_first:
      layer_output = source_attention_output + layer_output
    else:
      layer_output = self._output_layer_norm(layer_output + attention_output)

    return layer_output
//  [python]
def call(self, inputs):
    if isinstance(inputs, (list, tuple)) and len(inputs) == 2:
      input_tensor, attention_mask = inputs
    else:
      input_tensor, attention_mask = (inputs, None)

    if self._output_range:
      target_tensor = input_tensor[:, 0:self._output_range, :]
      attention_mask = attention_mask[:, 0:self._output_range, :]
    else:
      if self._norm_first:
        source_tensor = input_tensor
        input_tensor = self._attention_layer_norm(input_tensor)
      target_tensor = input_tensor

    attention_output = self._attention_layer(
        query=target_tensor, value=input_tensor, attention_mask=attention_mask)
    attention_output = self._attention_dropout(attention_output)
    if self._norm_first:
      attention_output = source_tensor + attention_output
    else:
      attention_output = self._attention_layer_norm(target_tensor +
                                                    attention_output)
    if self._norm_first:
      source_attention_output = attention_output
      attention_output = self._output_layer_norm(attention_output)

    layer_output = self._output_dense(attention_output)
    layer_output = self._output_dropout(layer_output)
    # During mixed precision training, attention_output is from layer norm and
    # is always fp32 for now. Cast layer_output to fp32 for the subsequent
    # add.
    layer_output = tf.cast(layer_output, tf.float32)
    if self._norm_first:
      layer_output = source_attention_output + layer_output
    else:
      layer_output = self._output_layer_norm(layer_output + attention_output)

    return layer_output
//  [python]
def transform_targets_for_output(y_true, grid_size, anchor_idxs):
    # y_true: (N, boxes, (x1, y1, x2, y2, class, best_anchor))
    N = tf.shape(y_true)[0]

    # y_true_out: (N, grid, grid, anchors, [x, y, w, h, obj, class])
    y_true_out = tf.zeros(
        (N, grid_size, grid_size, tf.shape(anchor_idxs)[0], 6))

    anchor_idxs = tf.to_int32(anchor_idxs)

    indexes = tf.TensorArray(tf.int32, 1, dynamic_size=True)
    updates = tf.TensorArray(tf.float32, 1, dynamic_size=True)
    idx = 0
    for i in tf.range(N):
        for j in tf.range(tf.shape(y_true)[1]):
            if tf.equal(y_true[i][j][2], 0):
                continue
            anchor_eq = tf.equal(
                anchor_idxs, tf.to_int32(y_true[i][j][5]))

            if tf.reduce_any(anchor_eq):
                box = y_true[i][j][0:4]
                box_xy = (y_true[i][j][0:2] + y_true[i][j][2:4]) / 2

                anchor_idx = tf.to_int32(tf.where(anchor_eq))
                grid_xy = tf.to_int32(box_xy // (1/grid_size))

                # grid[y][x][anchor] = (tx, ty, bw, bh, obj, class)
                indexes = indexes.write(
                    idx, [i, grid_xy[1], grid_xy[0], anchor_idx[0][0]])
                updates = updates.write(
                    idx, [box[0], box[1], box[2], box[3], 1, y_true[i][j][4]])
                idx += 1

    # tf.print(indexes.stack())
    # tf.print(updates.stack())

    return tf.tensor_scatter_nd_update(
        y_true_out, indexes.stack(), updates.stack())
//  [python]
def transform_targets_for_output(y_true, grid_size, anchor_idxs):
    # y_true: (N, boxes, (x1, y1, x2, y2, class, best_anchor))
    N = tf.shape(y_true)[0]

    # y_true_out: (N, grid, grid, anchors, [x, y, w, h, obj, class])
    y_true_out = tf.zeros(
        (N, grid_size, grid_size, tf.shape(anchor_idxs)[0], 6))

    anchor_idxs = tf.cast(anchor_idxs, tf.int32)

    indexes = tf.TensorArray(tf.int32, 1, dynamic_size=True)
    updates = tf.TensorArray(tf.float32, 1, dynamic_size=True)
    idx = 0
    for i in tf.range(N):
        for j in tf.range(tf.shape(y_true)[1]):
            if tf.equal(y_true[i][j][2], 0):
                continue
            anchor_eq = tf.equal(
                anchor_idxs, tf.cast(y_true[i][j][5], tf.int32))

            if tf.reduce_any(anchor_eq):
                box = y_true[i][j][0:4]
                box_xy = (y_true[i][j][0:2] + y_true[i][j][2:4]) / 2

                anchor_idx = tf.cast(tf.where(anchor_eq), tf.int32)
                grid_xy = tf.cast(box_xy // (1/grid_size), tf.int32)

                # grid[y][x][anchor] = (tx, ty, bw, bh, obj, class)
                indexes = indexes.write(
                    idx, [i, grid_xy[1], grid_xy[0], anchor_idx[0][0]])
                updates = updates.write(
                    idx, [box[0], box[1], box[2], box[3], 1, y_true[i][j][4]])
                idx += 1

    # tf.print(indexes.stack())
    # tf.print(updates.stack())

    return tf.tensor_scatter_nd_update(
        y_true_out, indexes.stack(), updates.stack())
//  [python]
def padding_mask(self, QK):
        padding = tf.to_float(tf.equal(QK,0))
        padding *= self.masking_num
        return QK+padding
//  [python]
def padding_mask(self, QK):
        padding = tf.cast(tf.equal(QK,0),tf.float32)
        padding *= self.masking_num
        return QK+padding
//  [python]
def reweight(msa1hot, cutoff):
    """reweight MSA based on cutoff"""
    with tf.name_scope('reweight'):
        id_min = tf.cast(tf.shape(msa1hot)[1], tf.float32) * cutoff
        id_mtx = tf.tensordot(msa1hot, msa1hot, [[1,2], [1,2]])
        id_mask = id_mtx > id_min
        w = 1.0/tf.reduce_sum(t.cast(id_mask, dtype=tf.float32),-1)
    return w
//  [python]
def reweight(msa1hot, cutoff):
    """reweight MSA based on cutoff"""
    with tf.name_scope('reweight'):
        id_min = tf.cast(tf.shape(msa1hot)[1], tf.float32) * cutoff
        id_mtx = tf.tensordot(msa1hot, msa1hot, [[1,2], [1,2]])
        id_mask = id_mtx > id_min
        w = 1.0/tf.reduce_sum(tf.cast(id_mask, dtype=tf.float32),-1)
    return w
//  [python]
def _compute_kernel(self, length: int):
    self._kernel = None
    # D x N x 1
    damping_factor, previous_timestep_weight = self._calc_coeffs()
    # D x N x L
    vander = t.cast(
        t.reshape(t.range(length), shape=(1, 1, length)),
        dtype=damping_factor.dtype,
    ) * t.math.log(previous_timestep_weight)
    kernel = (damping_factor * self.ema_expansion_matrix) * t.math.exp(vander)
    # D x L
    return t.einsum(
        "dnl,dn->dl", kernel, self.kernel_projection_matrix * self.scale
    )
//  [python]
def _compute_kernel(self, length: int):
    self._kernel = None
    # D x N x 1
    damping_factor, previous_timestep_weight = self._calc_coeffs()
    # D x N x L
    vander = tf.cast(
        tf.reshape(tf.range(length), shape=(1, 1, length)),
        dtype=damping_factor.dtype,
    ) * tf.math.log(previous_timestep_weight)
    kernel = (damping_factor * self.ema_expansion_matrix) * tf.math.exp(vander)
    # D x L
    return tf.einsum(
        "dnl,dn->dl", kernel, self.kernel_projection_matrix * self.scale
    )
//  [python]
def compute_proximity_loss(self):
        """Compute the second part (distance from x1) of the loss function."""
        proximity_loss = 0.0
        for i in range(self.total_CFs):
            proximity_loss += self.compute_dist(self.cfs[i], self.x1)
        return proximity_loss/t.cast((tf.multiply(len(self.minx[0]), self.total_CFs)), dtype=t.float32)
//  [python]
def compute_proximity_loss(self):
        """Compute the second part (distance from x1) of the loss function."""
        proximity_loss = 0.0
        for i in range(self.total_CFs):
            proximity_loss += self.compute_dist(self.cfs[i], self.x1)
        return proximity_loss/tf.cast((tf.multiply(len(self.minx[0]), self.total_CFs)), dtype=tf.float32)
//  [python]
def call(self, inputs):
        assert len(inputs) == 3, "inputs should be set [queries, keys, values]."
        queries, keys, values = inputs
        # 转换为32位
        if K.dtype(queries) != 'float32':  queries = K.cast(queries, 'float32')
        if K.dtype(keys) != 'float32':  keys = K.cast(keys, 'float32')
        if K.dtype(values) != 'float32':  values = K.cast(values, 'float32')
        # Qk计算
        matmul = tf.matmul(queries,keys,transpose_b=True)
        dk = tf.cast(tf.shape(keys)[-1],tf.float32)
        matmul = matmul / tf.sqrt(dk) # QxK后缩放dk**(0.5)
        # mask层,区别encoder和decoder部分
        if self.mode == "encoder":
            matmul = self.padding_mask(matmul)
        else:
            matmul = self.sequence_mask(matmul)
        softmax_out = K.softmax(matmul)  # SoftMax层
        return K.batch_dot(softmax_out, values)
//  [python]
def _prepare_image_and_label(self, data):
    """Prepare normalized image and label."""
    image = tf.io.decode_image(data['image/encoded'], channels=3)
    label = tf.io.decode_image(data['image/segmentation/class/encoded'],
                               channels=1)
    height = data['image/height']
    width = data['image/width']
    image = tf.reshape(image, (height, width, 3))

    label = tf.reshape(label, (1, height, width))
    label = tf.where(
        tf.math.greater(label, self.max_class), tf.zeros_like(label), label)
    label = tf.where(tf.math.equal(label, 0), tf.ones_like(label)*255, label)
    label = tf.to_float(label)
    # Normalizes image with mean and std pixel values.
    image = preprocess_ops.normalize_image(
        image, offset=[0.5, 0.5, 0.5], scale=[0.5, 0.5, 0.5])
    return image, label
//  [python]
def _prepare_image_and_label(self, data):
    """Prepare normalized image and label."""
    image = tf.io.decode_image(data['image/encoded'], channels=3)
    label = tf.io.decode_image(data['image/segmentation/class/encoded'],
                               channels=1)
    height = data['image/height']
    width = data['image/width']
    image = tf.reshape(image, (height, width, 3))

    label = tf.reshape(label, (1, height, width))
    label = tf.where(
        tf.math.greater(label, self.max_class), tf.zeros_like(label), label)
    label = tf.where(tf.math.equal(label, 0), tf.ones_like(label)*255, label)
    label = tf.cast(label, tf.float32)
    # Normalizes image with mean and std pixel values.
    image = preprocess_ops.normalize_image(
        image, offset=[0.5, 0.5, 0.5], scale=[0.5, 0.5, 0.5])
    return image, label
//  [python]
def build_losses(self, labels, model_outputs, aux_losses=None) -> tf.Tensor:
    label_ids_is_entity = labels[self.label_field_is_entity]
    label_ids_entity_type = labels[self.label_field_entity_type]
    loss_is_entity = tf_keras.losses.sparse_categorical_crossentropy(
        label_ids_is_entity,
        t.cast(model_outputs[0], tf.float32),
        from_logits=True)
    loss_entity_type = tf_keras.losses.sparse_categorical_crossentropy(
        label_ids_entity_type,
        t.cast(model_outputs[1], tf.float32),
        from_logits=True)
    loss = loss_is_entity + loss_entity_type

    if aux_losses:
      loss += tf.add_n(aux_losses)
    return tf_utils.safe_mean(loss)
//  [python]
def build_losses(self, labels, model_outputs, aux_losses=None) -> tf.Tensor:
    label_ids_is_entity = labels[self.label_field_is_entity]
    label_ids_entity_type = labels[self.label_field_entity_type]
    loss_is_entity = tf_keras.losses.sparse_categorical_crossentropy(
        label_ids_is_entity,
        tf.cast(model_outputs[0], tf.float32),
        from_logits=True)
    loss_entity_type = tf_keras.losses.sparse_categorical_crossentropy(
        label_ids_entity_type,
        tf.cast(model_outputs[1], tf.float32),
        from_logits=True)
    loss = loss_is_entity + loss_entity_type

    if aux_losses:
      loss += tf.add_n(aux_losses)
    return tf_utils.safe_mean(loss)
//  [python]
def validation_step(self,
                      inputs: Tuple[Any, Any],
                      model: tf_keras.Model,
                      metrics: Optional[List[Any]] = None):
    """Validation step.

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

    Returns:
      A dictionary of logs.
    """
    features, labels = inputs

    outputs = model(features, training=False)
    outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs)

    losses = self.build_losses(outputs['raw_output'], labels)

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

    coco_model_outputs = {
        'detection_boxes': outputs['boxes'],
        'detection_scores': outputs['confidence'],
        'detection_classes': outputs['classes'],
        'num_detections': outputs['num_detections'],
        'source_id': labels['groundtruths']['source_id'],
        'image_info': labels['image_info']
    }

    logs.update({self.coco_metric.name: (labels['groundtruths'],
                                         coco_model_outputs)})

    if metrics:
      for m in metrics:
        m.update_state(losses[m.name])
        logs.update({m.name: m.result()})
    return logs
//  [python]
def call(self, x, padding_mask: Optional[tf.Tensor] = None) -> tf.Tensor:
    """Input shape: Time x Batch x Channel.

    Args:
      x: Tensor input.
      padding_mask (ByteTensor, optional): mask to exclude keys that are pads,
        of shape `(batch, src_len)`, where padding elements are indicated by
        1s.
    Returns:
      transformed: transformed Tensor.
    """

    seq_len, _, embed_dim = x.shape
    assert embed_dim == self.embed_dim
    if seq_len is None:
      seq_len = 1

    # L x B x D
    residual = x * self.residual_weight

    # L x B x D -> B x D x L
    x = tf.transpose(x, perm=(1, 2, 0))

    # Masking of the tensor
    if padding_mask is not None:
      x = x * tf.cast(tf.expand_dims(padding_mask, axis=1), x.dtype)

    k = self.kernel(seq_len)

    kernel_size = k.shape[1]
    fft_len = seq_len
    s = 0

    if self.bidirectional:
      k1, k2 = tf.split(k, [self.embed_dim, self.embed_dim], axis=0)
      # D x 2*L-1
      padding_l = tf.constant([[0, 0], [kernel_size - 1, 0]])
      padding_r = tf.constant([[0, 0], [0, kernel_size - 1]])
      padding_x = tf.constant([[0, 0], [0, 0], [kernel_size - 1, 0]])
      k = tf.pad(k1, padding_l) + tf.pad(tf.reverse(k2, axis=[-1]), padding_r)
      x = tf.pad(x, padding_x)
      fft_len = fft_len + kernel_size - 1
      s = 2 * kernel_size - 2

    k_f = tf.signal.rfft(
        k, fft_length=tf.constant([2 * fft_len], dtype=tf.int32)
    )
    x_f = tf.signal.rfft(
        x, fft_length=tf.constant([2 * fft_len], dtype=tf.int32)
    )
    # B x D x L
    out = tf.signal.irfft(
        x_f * k_f, fft_length=tf.constant([2 * fft_len], dtype=tf.int32)
    )[..., s : s + seq_len]

    # B x D x L -> L x B x D
    out = tf.nn.silu(tf.transpose(out, perm=(2, 0, 1)) + residual)
    return out
//  [python]
def train_step(self,
                 inputs: Tuple[Any, Any],
                 model: tf_keras.Model,
                 optimizer: tf_keras.optimizers.Optimizer,
                 metrics: Optional[List[Any]] = None):
    """Does forward and backward.

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

    Returns:
      A dictionary of logs.
    """
    features, labels = inputs

    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
    with tf.GradientTape() as tape:
      outputs = model(features, training=True)
      # Casting output layer as float32 is necessary when mixed_precision is
      # mixed_float16 or mixed_bfloat16 to ensure output is casted as float32.
      outputs = tf.nest.map_structure(lambda x: tf.to_float(x), outputs)

      losses = self.build_losses(outputs['raw_output'], labels)

      scaled_loss = losses['total_loss'] / num_replicas
      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
      # scaled for numerical stability.
      if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    # compute the gradient
    tvars = model.trainable_variables
    gradients = tape.gradient(scaled_loss, tvars)

    # get unscaled loss if the scaled loss was used
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      gradients = optimizer.get_unscaled_gradients(gradients)

    if self.task_config.gradient_clip_norm > 0.0:
      gradients, _ = tf.clip_by_global_norm(gradients,
                                            self.task_config.gradient_clip_norm)

    optimizer.apply_gradients(list(zip(gradients, tvars)))

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

    if metrics:
      for m in metrics:
        m.update_state(losses[m.name])
        logs.update({m.name: m.result()})

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

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

    Returns:
      A dictionary of logs.
    """
    features, labels = inputs

    num_replicas = tf.distribute.get_strategy().num_replicas_in_sync
    with tf.GradientTape() as tape:
      outputs = model(features, training=True)
      # Casting output layer as float32 is necessary when mixed_precision is
      # mixed_float16 or mixed_bfloat16 to ensure output is casted as float32.
      outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs)

      losses = self.build_losses(outputs['raw_output'], labels)

      scaled_loss = losses['total_loss'] / num_replicas
      # For mixed_precision policy, when LossScaleOptimizer is used, loss is
      # scaled for numerical stability.
      if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
        scaled_loss = optimizer.get_scaled_loss(scaled_loss)

    # compute the gradient
    tvars = model.trainable_variables
    gradients = tape.gradient(scaled_loss, tvars)

    # get unscaled loss if the scaled loss was used
    if isinstance(optimizer, tf_keras.mixed_precision.LossScaleOptimizer):
      gradients = optimizer.get_unscaled_gradients(gradients)

    if self.task_config.gradient_clip_norm > 0.0:
      gradients, _ = tf.clip_by_global_norm(gradients,
                                            self.task_config.gradient_clip_norm)

    optimizer.apply_gradients(list(zip(gradients, tvars)))

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

    if metrics:
      for m in metrics:
        m.update_state(losses[m.name])
        logs.update({m.name: m.result()})

    return logs
//  [python]
def compute_loss(
        model_output: tf.Tensor,
        gt: tf.Tensor,
        seq_len: tf.Tensor,
    ) -> tf.Tensor:
        """Compute categorical cross-entropy loss for the model.
        Sequences are masked after the EOS character.

        Args:
        ----
            gt: the encoded tensor with gt labels
            model_output: predicted logits of the model
            seq_len: lengths of each gt word inside the batch

        Returns:
        -------
            The loss of the model on the batch
        """
        # Input length : number of timesteps
        input_len = tf.shape(model_output)[1]
        # Add one for additional <eos> token
        seq_len = seq_len + 1
        # One-hot gt labels
        oh_gt = tf.one_hot(gt, depth=model_output.shape[2])
        # Compute loss
        cce = tf.nn.softmax_cross_entropy_with_logits(oh_gt, model_output)
        # Compute mask
        mask_values = tf.zeros_like(cce)
        mask_2d = tf.sequence_mask(seq_len, input_len)
        masked_loss = tf.where(mask_2d, cce, mask_values)
        ce_loss = tf.math.divide(tf.reduce_sum(masked_loss, axis=1), tf.to_float(seq_len))
        return tf.expand_dims(ce_loss, axis=1)
//  [python]
def compute_loss(
        model_output: tf.Tensor,
        gt: tf.Tensor,
        seq_len: tf.Tensor,
    ) -> tf.Tensor:
        """Compute categorical cross-entropy loss for the model.
        Sequences are masked after the EOS character.

        Args:
        ----
            gt: the encoded tensor with gt labels
            model_output: predicted logits of the model
            seq_len: lengths of each gt word inside the batch

        Returns:
        -------
            The loss of the model on the batch
        """
        # Input length : number of timesteps
        input_len = tf.shape(model_output)[1]
        # Add one for additional <eos> token
        seq_len = seq_len + 1
        # One-hot gt labels
        oh_gt = tf.one_hot(gt, depth=model_output.shape[2])
        # Compute loss
        cce = tf.nn.softmax_cross_entropy_with_logits(oh_gt, model_output)
        # Compute mask
        mask_values = tf.zeros_like(cce)
        mask_2d = tf.sequence_mask(seq_len, input_len)
        masked_loss = tf.where(mask_2d, cce, mask_values)
        ce_loss = tf.math.divide(tf.reduce_sum(masked_loss, axis=1), tf.cast(seq_len, model_output.dtype))
        return tf.expand_dims(ce_loss, axis=1)