jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def scaled_dot_product_attention(q, k, v, mask, is_training=True, rel_key_emb=None, rel_val_emb=None):
        # attn_dim: num_joints for spatial and seq_len for temporal
        '''
        The scaled dot product attention mechanism introduced in the Transformer
        :param q: the query vectors matrix (..., attn_dim, d_model/num_heads)
        :param k: the key vector matrix (..., attn_dim, d_model/num_heads)
        :param v: the value vector matrix (..., attn_dim, d_model/num_heads)
        :param mask: a mask for attention
        :return: the updated encoding and the attention weights matrix
        '''

        matmul_qk = tf.matmul(q, k, transpose_b=True)  # (..., num_heads, attn_dim, attn_dim)

        batch_size = tf.shape(q)[0]
        heads = tf.shape(q)[1]
        length = tf.shape(q)[2]
        
        if rel_key_emb is not None:
            q_t = tf.transpose(q, [2, 0, 1, 3])
            q_t_r = tf.reshape(q_t, [length, heads*batch_size, -1])
            q_tz_matmul = tf.matmul(q_t_r, rel_key_emb, transpose_b=True)
            q_tz_matmul_r = tf.reshape(q_tz_matmul, [length, batch_size, heads, -1])
            q_tz_matmul_r_t = tf.transpose(q_tz_matmul_r, [1, 2, 0, 3])
            matmul_qk += q_tz_matmul_r_t

        # scale matmul_qk
        dk = tf.cast(tf.shape(k)[-1], tf.float32)
        scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)

        # add the mask to the scaled tensor.
        if mask is not None:
            scaled_attention_logits += (mask * -1e9)

        # normalized on the last axis (seq_len_k) so that the scores add up to 1.
        attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)  # (..., num_heads, attn_dim, attn_dim)

        # attention_weights = tf.layers.dropout(attention_weights, training=is_training, rate=0.2)
        output = tf.matmul(attention_weights, v)  # (..., num_heads, attn_dim, depth)
        
        if rel_val_emb is not None:
            w_t = tf.transpose(attention_weights, [2, 0, 1, 3])
            w_t_r = tf.reshape(w_t, [length, heads*batch_size, -1])
            w_tz_matmul = tf.matmul(w_t_r, rel_val_emb, transpose_b=False)
            w_tz_matmul_r = tf.reshape(w_tz_matmul, [length, batch_size, heads, -1])
            w_tz_matmul_r_t = tf.transpose(w_tz_matmul_r, [1, 2, 0, 3])
            output += w_tz_matmul_r_t

        return output, attention_weights
//  [python]
def validation_step(
      self,
      inputs,
      model: tf_keras.Model,
      metrics: Optional[Sequence[tf_keras.metrics.Metric]] = None
  ) -> Dict[Any, Any]:
    """Validatation step.

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

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

    input_partition_dims = self.task_config.eval_input_partition_dims
    if input_partition_dims:
      strategy = tf.distribute.get_strategy()
      features = strategy.experimental_split_to_logical_devices(
          features, input_partition_dims)

    outputs = self.inference_step(features, model)
    outputs = tf.nest.map_structure(lambda x: tf.cast(x, tf.float32), outputs)
    outputs = outputs['logits']
    if self.task_config.model.head.output_logits:
      outputs = tf.nn.softmax(outputs)

    loss = self.build_losses(
        model_outputs=outputs, labels=labels, aux_losses=model.losses)
    logs = {self.loss: loss}

    # Compute dice score metrics on CPU.
    for metric in self.metrics:
      labels = tf.cast(labels, tf.float32)
      logits = tf.cast(outputs, tf.float32)
      logs.update({metric.name: (labels, logits)})

    return logs
//  [python]
def aggregate_logs(
      self,
      state: Optional[Sequence[Union[segmentation_metrics.DiceScore,
                                     tf_keras.metrics.Metric]]] = None,
      step_outputs: Optional[Mapping[str, Any]] = None
  ) -> Sequence[tf_keras.metrics.Metric]:
    """Aggregates statistics to compute metrics over training.

    Args:
      state: A sequence of tf_keras.metrics.Metric objects. Each element records
        a metric.
      step_outputs: A dictionary of [metric_name, (labels, output)] from a step.

    Returns:
      An updated sequence of tf_keras.metrics.Metric objects.
    """
    if state is None:
      for metric in self.metrics:
        metric.reset_states()
      state = self.metrics

    for metric in self.metrics:
      labels = step_outputs[metric.name][0]
      predictions = step_outputs[metric.name][1]

      # If `step_output` is distributed, it contains a tuple of Tensors instead
      # of a single Tensor, so we need to concatenate them along the batch
      # dimension in this case to have a single Tensor.
      if isinstance(labels, tuple):
        labels = tf.concat(list(labels), axis=0)
      if isinstance(predictions, tuple):
        predictions = tf.concat(list(predictions), axis=0)

      labels = t.cast(labels, t.float32)
      predictions = t.cast(predictions, t.float32)
      metric.update_state(labels, predictions)
    return state
//  [python]
def aggregate_logs(
      self,
      state: Optional[Sequence[Union[segmentation_metrics.DiceScore,
                                     tf_keras.metrics.Metric]]] = None,
      step_outputs: Optional[Mapping[str, Any]] = None
  ) -> Sequence[tf_keras.metrics.Metric]:
    """Aggregates statistics to compute metrics over training.

    Args:
      state: A sequence of tf_keras.metrics.Metric objects. Each element records
        a metric.
      step_outputs: A dictionary of [metric_name, (labels, output)] from a step.

    Returns:
      An updated sequence of tf_keras.metrics.Metric objects.
    """
    if state is None:
      for metric in self.metrics:
        metric.reset_states()
      state = self.metrics

    for metric in self.metrics:
      labels = step_outputs[metric.name][0]
      predictions = step_outputs[metric.name][1]

      # If `step_output` is distributed, it contains a tuple of Tensors instead
      # of a single Tensor, so we need to concatenate them along the batch
      # dimension in this case to have a single Tensor.
      if isinstance(labels, tuple):
        labels = tf.concat(list(labels), axis=0)
      if isinstance(predictions, tuple):
        predictions = tf.concat(list(predictions), axis=0)

      labels = tf.cast(labels, tf.float32)
      predictions = tf.cast(predictions, tf.float32)
      metric.update_state(labels, predictions)
    return state
//  [python]
def _compute_loss(self, prediction_tensor, target_tensor, weights):
    """Compute loss function.

    Args:
      prediction_tensor: A float tensor of shape [batch_size, num_anchors,
        num_classes] representing the predicted logits for each class
      target_tensor: A float tensor of shape [batch_size, num_anchors,
        num_classes] representing one-hot encoded classification targets
      weights: a float tensor of shape [batch_size, num_anchors]

    Returns:
      loss: a (scalar) tensor representing the value of the loss function
            or a float tensor of shape [batch_size, num_anchors]
    """
    if self._bootstrap_type == 'soft':
      bootstrap_target_tensor = self._alpha * target_tensor + (
          1.0 - self._alpha) * tf.sigmoid(prediction_tensor)
    else:
      bootstrap_target_tensor = self._alpha * target_tensor + (
          1.0 - self._alpha) * tf.to_float(
              tf.sigmoid(prediction_tensor) > 0.5)
    per_entry_cross_ent = (tf.nn.sigmoid_cross_entropy_with_logits(
        labels=bootstrap_target_tensor, logits=prediction_tensor))
    if self._anchorwise_output:
      return tf.reduce_sum(per_entry_cross_ent * tf.expand_dims(weights, 2), 2)
    return tf.reduce_sum(per_entry_cross_ent * tf.expand_dims(weights, 2))
//  [python]
def _compute_loss(self, prediction_tensor, target_tensor, weights):
    """Compute loss function.

    Args:
      prediction_tensor: A float tensor of shape [batch_size, num_anchors,
        num_classes] representing the predicted logits for each class
      target_tensor: A float tensor of shape [batch_size, num_anchors,
        num_classes] representing one-hot encoded classification targets
      weights: a float tensor of shape [batch_size, num_anchors]

    Returns:
      loss: a (scalar) tensor representing the value of the loss function
            or a float tensor of shape [batch_size, num_anchors]
    """
    if self._bootstrap_type == 'soft':
      bootstrap_target_tensor = self._alpha * target_tensor + (
          1.0 - self._alpha) * tf.sigmoid(prediction_tensor)
    else:
      bootstrap_target_tensor = self._alpha * target_tensor + (
          1.0 - self._alpha) * tf.cast(
              tf.sigmoid(prediction_tensor) > 0.5, tf.float32)
    per_entry_cross_ent = (tf.nn.sigmoid_cross_entropy_with_logits(
        labels=bootstrap_target_tensor, logits=prediction_tensor))
    if self._anchorwise_output:
      return tf.reduce_sum(per_entry_cross_ent * tf.expand_dims(weights, 2), 2)
    return tf.reduce_sum(per_entry_cross_ent * tf.expand_dims(weights, 2))
//  [python]
def ensure_tensor(inputs, dtype=None):
    """Ensures the input is a Tensor, SparseTensor or RaggedTensor."""
    if not isinstance(inputs, (tf.Tensor, tf.RaggedTensor, tf.SparseTensor)):
        inputs = tf.convert_to_tensor(inputs, dtype)
    if dtype is not None and inputs.dtype != dtype:
        inputs = tf.cast(inputs, dtype)
    return inputs
//  [python]
def cast_to(inputs, dtype):
    if IMAGES in inputs:
        inputs[IMAGES] = t.cast(inputs[IMAGES], dtype)
    if LABELS in inputs:
        inputs[LABELS] = t.cast(inputs[LABELS], dtype)
    if BOUNDING_BOXES in inputs:
        inputs[BOUNDING_BOXES]["boxes"] = t.cast(
            inputs[BOUNDING_BOXES]["boxes"], dtype
        )
        inputs[BOUNDING_BOXES]["classes"] = t.cast(
            inputs[BOUNDING_BOXES]["classes"], dtype
        )
    if SEGMENTATION_MASKS in inputs:
        inputs[SEGMENTATION_MASKS] = t.cast(inputs[SEGMENTATION_MASKS], dtype)
    if KEYPOINTS in inputs:
        inputs[KEYPOINTS] = t.cast(inputs[KEYPOINTS], dtype)
    if CUSTOM_ANNOTATIONS in inputs:
        raise NotImplementedError()
    return inputs
//  [python]
def cast_to(inputs, dtype):
    if IMAGES in inputs:
        inputs[IMAGES] = tf.cast(inputs[IMAGES], dtype)
    if LABELS in inputs:
        inputs[LABELS] = tf.cast(inputs[LABELS], dtype)
    if BOUNDING_BOXES in inputs:
        inputs[BOUNDING_BOXES]["boxes"] = tf.cast(
            inputs[BOUNDING_BOXES]["boxes"], dtype
        )
        inputs[BOUNDING_BOXES]["classes"] = tf.cast(
            inputs[BOUNDING_BOXES]["classes"], dtype
        )
    if SEGMENTATION_MASKS in inputs:
        inputs[SEGMENTATION_MASKS] = tf.cast(inputs[SEGMENTATION_MASKS], dtype)
    if KEYPOINTS in inputs:
        inputs[KEYPOINTS] = tf.cast(inputs[KEYPOINTS], dtype)
    if CUSTOM_ANNOTATIONS in inputs:
        raise NotImplementedError()
    return inputs
//  [python]
def visualize_images(name, *imgs):
    """Generate tensor for TensorBoard (casting, clipping)

    Args:
        name: name for visualiaztion operation
        *imgs: multiple images in *args style

    Example:

        visualize_images('viz1', img1)
        visualize_images('viz2', img1, img2, img3)
    """
    xy = (tf.concat(imgs, axis=2) + 1.) * 128.
    xy = tf.clip_by_value(xy, 0, 255)
    xy = tf.cast(xy, tf.uint8, name='viz')
    tf.summary.image(name, xy, max_outputs=30)
//  [python]
def visualize_images(name, *imgs):
    """Generate tensor for TensorBoard (casting, clipping)

    Args:
        name: name for visualiaztion operation
        *imgs: multiple images in *args style

    Example:

        visualize_images('viz1', img1)
        visualize_images('viz2', img1, img2, img3)
    """
    xy = (tf.concat(imgs, axis=2) + 1.) * 128.
    xy = tf.cast(tf.clip_by_value(xy, 0, 255), tf.uint8, name='viz')
    tf.summary.image(name, xy, max_outputs=30)
//  [python]
def train_mnist(steps_per_epoch=100, epochs=5, 
                lr=1e-3, layers=2, hidden_size=64, 
                logger_kwargs=dict(), save_freq=1):

    logger = EpochLogger(**logger_kwargs)
    logger.save_config(locals())

    # Load and preprocess MNIST data
    (x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
    x_train = x_train.reshape(-1, 28*28) / 255.0

    # Define inputs & main outputs from computation graph
    x_ph = tf.placeholder(tf.float32, shape=(None, 28*28))
    y_ph = tf.placeholder(tf.int32, shape=(None,))
    logits = mlp(x_ph, hidden_sizes=[hidden_size]*layers + [10], activation=tf.nn.relu)
    predict = tf.argmax(logits, axis=1, output_type=tf.int32)

    # Define loss function, accuracy, and training op
    y = tf.one_hot(y_ph, 10)
    loss = tf.losses.softmax_cross_entropy(y, logits)
    acc = tf.reduce_mean(tf.to_float(tf.equal(y_ph, predict)))
    train_op = tf.train.AdamOptimizer().minimize(loss)

    # Prepare session
    sess = tf.Session()
    sess.run(tf.global_variables_initializer())

    # Setup model saving
    logger.setup_tf_saver(sess, inputs={'x': x_ph}, 
                                outputs={'logits': logits, 'predict': predict})

    start_time = time.time()

    # Run main training loop
    for epoch in range(epochs):
        for t in range(steps_per_epoch):
            idxs = np.random.randint(0, len(x_train), 32)
            feed_dict = {x_ph: x_train[idxs],
                         y_ph: y_train[idxs]}
            outs = sess.run([loss, acc, train_op], feed_dict=feed_dict)
            logger.store(Loss=outs[0], Acc=outs[1])

        # Save model
        if (epoch % save_freq == 0) or (epoch == epochs-1):
            logger.save_state(state_dict=dict(), itr=None)

        # Log info about epoch
        logger.log_tabular('Epoch', epoch)
        logger.log_tabular('Acc', with_min_and_max=True)
        logger.log_tabular('Loss', average_only=True)
        logger.log_tabular('TotalGradientSteps', (epoch+1)*steps_per_epoch)
        logger.log_tabular('Time', time.time()-start_time)
        logger.dump_tabular()
//  [python]
def train_mnist(steps_per_epoch=100, epochs=5, 
                lr=1e-3, layers=2, hidden_size=64, 
                logger_kwargs=dict(), save_freq=1):

    logger = EpochLogger(**logger_kwargs)
    logger.save_config(locals())

    # Load and preprocess MNIST data
    (x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
    x_train = x_train.reshape(-1, 28*28) / 255.0

    # Define inputs & main outputs from computation graph
    x_ph = tf.placeholder(tf.float32, shape=(None, 28*28))
    y_ph = tf.placeholder(tf.int32, shape=(None,))
    logits = mlp(x_ph, hidden_sizes=[hidden_size]*layers + [10], activation=tf.nn.relu)
    predict = tf.argmax(logits, axis=1, output_type=tf.int32)

    # Define loss function, accuracy, and training op
    y = tf.one_hot(y_ph, 10)
    loss = tf.losses.softmax_cross_entropy(y, logits)
    acc = tf.reduce_mean(tf.cast(tf.equal(y_ph, predict), tf.float32))
    train_op = tf.train.AdamOptimizer().minimize(loss)

    # Prepare session
    sess = tf.Session()
    sess.run(tf.global_variables_initializer())

    # Setup model saving
    logger.setup_tf_saver(sess, inputs={'x': x_ph}, 
                                outputs={'logits': logits, 'predict': predict})

    start_time = time.time()

    # Run main training loop
    for epoch in range(epochs):
        for t in range(steps_per_epoch):
            idxs = np.random.randint(0, len(x_train), 32)
            feed_dict = {x_ph: x_train[idxs],
                         y_ph: y_train[idxs]}
            outs = sess.run([loss, acc, train_op], feed_dict=feed_dict)
            logger.store(Loss=outs[0], Acc=outs[1])

        # Save model
        if (epoch % save_freq == 0) or (epoch == epochs-1):
            logger.save_state(state_dict=dict(), itr=None)

        # Log info about epoch
        logger.log_tabular('Epoch', epoch)
        logger.log_tabular('Acc', with_min_and_max=True)
        logger.log_tabular('Loss', average_only=True)
        logger.log_tabular('TotalGradientSteps', (epoch+1)*steps_per_epoch)
        logger.log_tabular('Time', time.time()-start_time)
        logger.dump_tabular()
//  [python]
def get_images_shape(images, dtype=tf.int32):
    """Get ``heights`` and ``widths`` of the input images.

    Input images can be ``tf.Tensor`` or ``tf.RaggedTensor`` with the shape of
    [B, H|None, W|None, C].

    Args:
        images (tf.Tensor|tf.RaggedTensor): The input images.
        dtype (tf.dtypes.DType, optional): The dtype of the outputs. Defaults to
            ``tf.int32``.
    """
    if isinstance(images, tf.RaggedTensor):
        heights = tf.reshape(images.row_lengths(), (-1, 1))
        widths = tf.reshape(
            tf.reduce_max(images.row_lengths(axis=2), 1), (-1, 1)
        )
    else:
        batch_size = tf.shape(images)[0]
        heights = tf.repeat(tf.shape(images)[H_AXIS], repeats=[batch_size])
        heights = tf.reshape(heights, shape=(-1, 1))
        widths = tf.repeat(tf.shape(images)[W_AXIS], repeats=[batch_size])
        widths = tf.reshape(widths, shape=(-1, 1))
    return tf.to_int32(heights), tf.to_int32(widths)
//  [python]
def get_images_shape(images, dtype=tf.int32):
    """Get ``heights`` and ``widths`` of the input images.

    Input images can be ``tf.Tensor`` or ``tf.RaggedTensor`` with the shape of
    [B, H|None, W|None, C].

    Args:
        images (tf.Tensor|tf.RaggedTensor): The input images.
        dtype (tf.dtypes.DType, optional): The dtype of the outputs. Defaults to
            ``tf.int32``.
    """
    if isinstance(images, tf.RaggedTensor):
        heights = tf.reshape(images.row_lengths(), (-1, 1))
        widths = tf.reshape(
            tf.reduce_max(images.row_lengths(axis=2), 1), (-1, 1)
        )
    else:
        batch_size = tf.shape(images)[0]
        heights = tf.repeat(tf.shape(images)[H_AXIS], repeats=[batch_size])
        heights = tf.reshape(heights, shape=(-1, 1))
        widths = tf.repeat(tf.shape(images)[W_AXIS], repeats=[batch_size])
        widths = tf.reshape(widths, shape=(-1, 1))
    return tf.cast(heights, dtype=dtype), tf.cast(widths, dtype=dtype)
//  [python]
def get_position_params(
    tops, bottoms, lefts, rights, position, random_generator
):
    """This function supposes arguments are at `center` padding method."""
    tops = tf.convert_to_tensor(tops)
    bottoms = tf.convert_to_tensor(bottoms)
    lefts = tf.convert_to_tensor(lefts)
    rights = tf.convert_to_tensor(rights)

    if position == PaddingPosition.CENTER:
        # do nothing
        bottoms = bottoms
        rights = rights
        tops = tops
        lefts = lefts
    elif position == PaddingPosition.TOP_LEFT:
        bottoms += tops
        rights += lefts
        tops = tf.zeros_like(tops)
        lefts = tf.zeros_like(lefts)
    elif position == PaddingPosition.TOP_RIGHT:
        bottoms += tops
        lefts += rights
        tops = tf.zeros_like(tops)
        rights = tf.zeros_like(rights)
    elif position == PaddingPosition.BOTTOM_LEFT:
        tops += bottoms
        rights += lefts
        bottoms = tf.zeros_like(bottoms)
        lefts = tf.zeros_like(lefts)
    elif position == PaddingPosition.BOTTOM_RIGHT:
        tops += bottoms
        lefts += rights
        bottoms = tf.zeros_like(bottoms)
        rights = tf.zeros_like(rights)
    elif position == PaddingPosition.RANDOM:
        batch_size = tf.shape(tops)[0]
        original_dtype = tops.dtype
        h_pads = tf.to_float(tops + bottoms)
        w_pads = tf.to_float(lefts + rights)
        tops = random_generator.random_uniform(
            shape=(batch_size, 1), minval=0, maxval=1, dtype=tf.float32
        )
        tops = tf.to_int32(tf.round(tops * h_pads))
        bottoms = tf.to_int32(h_pads) - tops
        lefts = random_generator.random_uniform(
            shape=(batch_size, 1), minval=0, maxval=1, dtype=tf.float32
        )
        lefts = tf.to_int32(tf.round(lefts * w_pads))
        rights = tf.to_int32(w_pads) - lefts
    else:
        raise NotImplementedError(
            f"Value not recognized for `position`: {position}. Supported "
            f"values are: {PaddingPosition._value2member_map_.keys()}"
        )

    return tops, bottoms, lefts, rights
//  [python]
def get_position_params(
    tops, bottoms, lefts, rights, position, random_generator
):
    """This function supposes arguments are at `center` padding method."""
    tops = tf.convert_to_tensor(tops)
    bottoms = tf.convert_to_tensor(bottoms)
    lefts = tf.convert_to_tensor(lefts)
    rights = tf.convert_to_tensor(rights)

    if position == PaddingPosition.CENTER:
        # do nothing
        bottoms = bottoms
        rights = rights
        tops = tops
        lefts = lefts
    elif position == PaddingPosition.TOP_LEFT:
        bottoms += tops
        rights += lefts
        tops = tf.zeros_like(tops)
        lefts = tf.zeros_like(lefts)
    elif position == PaddingPosition.TOP_RIGHT:
        bottoms += tops
        lefts += rights
        tops = tf.zeros_like(tops)
        rights = tf.zeros_like(rights)
    elif position == PaddingPosition.BOTTOM_LEFT:
        tops += bottoms
        rights += lefts
        bottoms = tf.zeros_like(bottoms)
        lefts = tf.zeros_like(lefts)
    elif position == PaddingPosition.BOTTOM_RIGHT:
        tops += bottoms
        lefts += rights
        bottoms = tf.zeros_like(bottoms)
        rights = tf.zeros_like(rights)
    elif position == PaddingPosition.RANDOM:
        batch_size = tf.shape(tops)[0]
        original_dtype = tops.dtype
        h_pads = tf.cast(tops + bottoms, dtype=tf.float32)
        w_pads = tf.cast(lefts + rights, dtype=tf.float32)
        tops = random_generator.random_uniform(
            shape=(batch_size, 1), minval=0, maxval=1, dtype=tf.float32
        )
        tops = tf.cast(tf.round(tops * h_pads), dtype=original_dtype)
        bottoms = tf.cast(h_pads, dtype=tf.int32) - tops
        lefts = random_generator.random_uniform(
            shape=(batch_size, 1), minval=0, maxval=1, dtype=tf.float32
        )
        lefts = tf.cast(tf.round(lefts * w_pads), dtype=original_dtype)
        rights = tf.cast(w_pads, dtype=tf.int32) - lefts
    else:
        raise NotImplementedError(
            f"Value not recognized for `position`: {position}. Supported "
            f"values are: {PaddingPosition._value2member_map_.keys()}"
        )

    return tops, bottoms, lefts, rights
//  [python]
def read_smpl_data(filename_queue):
    """
    Parses a smpl Example proto.
    It's contents are:
        'pose'  : 72-D float
        'shape' : 10-D float
    """
    with tf.name_scope(None, 'read_smpl_data', [filename_queue]):
        reader = tf.TFRecordReader()
        _, example_serialized = reader.read(filename_queue)

        feature_map = {
            'pose': tf.FixedLenFeature((72, ), dtype=tf.float32),
            'shape': tf.FixedLenFeature((10, ), dtype=tf.float32)
        }

        features = tf.parse_single_example(example_serialized, feature_map)
        pose = tf.to_float(features['pose'])
        shape = tf.to_float(features['shape'])

        return pose, shape
//  [python]
def read_smpl_data(filename_queue):
    """
    Parses a smpl Example proto.
    It's contents are:
        'pose'  : 72-D float
        'shape' : 10-D float
    """
    with tf.name_scope(None, 'read_smpl_data', [filename_queue]):
        reader = tf.TFRecordReader()
        _, example_serialized = reader.read(filename_queue)

        feature_map = {
            'pose': tf.FixedLenFeature((72, ), dtype=tf.float32),
            'shape': tf.FixedLenFeature((10, ), dtype=tf.float32)
        }

        features = tf.parse_single_example(example_serialized, feature_map)
        pose = tf.cast(features['pose'], dtype=tf.float32)
        shape = tf.cast(features['shape'], dtype=tf.float32)

        return pose, shape
//  [python]
def get_pseudo_labels(
    hard_labels,
    reg,
    bag_size,
    avg_true_label,
    preds,
):
  """Pseudo labels for ER-OT."""
  p = tf.stack([1 - avg_true_label, avg_true_label])
  b = tf.ones(bag_size, dtype=tf.float32) / tf.cast(bag_size, dtype=tf.float32)
  # pylint: disable=invalid-name
  P = tf.stack([1 - preds, preds])
  Q = sink(p, b, -tf.math.log(P), reg)
  pl = tf.cast(bag_size, dtype=tf.float32) * Q[1, :]
  if hard_labels:
    pl = (tf.math.sign(pl - 0.5) + 1.0) / 2.0
  return pl
//  [python]
def dllp_loss_graph(
    batch_size,
    instance_loss,
    x,
    y,
    y_pred,
):
  """DLLP loss for graph execution."""
  del x
  y_sg = tf.stop_gradient(y)
  loss = tf.zeros((), dtype=tf.float32)
  start = tf.zeros((), dtype=tf.int32)
  for _ in range(batch_size):
    bag_size = tf.to_int32(tf.gather(y_sg, start, axis=0)[0])
    avg_true_label = (
        tf.gather(y_sg, start + 1, axis=0)[0]
        / tf.gather(y_sg, start, axis=0)[0]
    )
    avg_pred_label = tf.reduce_mean(
        tf.gather(y_pred, tf.range(start, start + bag_size), axis=0)
    )
    loss += instance_loss(
        tf.convert_to_tensor([avg_true_label]),
        tf.convert_to_tensor([avg_pred_label]),
    )
    start += bag_size
  return loss / batch_size
//  [python]
def dllp_loss_graph(
    batch_size,
    instance_loss,
    x,
    y,
    y_pred,
):
  """DLLP loss for graph execution."""
  del x
  y_sg = tf.stop_gradient(y)
  loss = tf.zeros((), dtype=tf.float32)
  start = tf.zeros((), dtype=tf.int32)
  for _ in range(batch_size):
    bag_size = tf.cast(tf.gather(y_sg, start, axis=0)[0], dtype=tf.int32)
    avg_true_label = (
        tf.gather(y_sg, start + 1, axis=0)[0]
        / tf.gather(y_sg, start, axis=0)[0]
    )
    avg_pred_label = tf.reduce_mean(
        tf.gather(y_pred, tf.range(start, start + bag_size), axis=0)
    )
    loss += instance_loss(
        tf.convert_to_tensor([avg_true_label]),
        tf.convert_to_tensor([avg_pred_label]),
    )
    start += bag_size
  return loss / batch_size
//  [python]
def erot_loss(
    hard_labels,
    reg,
    batch_size,
    instance_loss,
    x,
    y,
    y_pred,
):
  """EROT loss."""
  del x
  y_sg = tf.stop_gradient(y)
  loss = tf.zeros((), dtype=tf.float32)
  start = tf.zeros((), dtype=tf.int32)
  for _ in range(batch_size):
    bag_size = tf.to_int32(tf.gather(y_sg, start, axis=0)[0])
    avg_true_label = (
        tf.gather(y_sg, start + 1, axis=0)[0]
        / tf.gather(y_sg, start, axis=0)[0]
    )
    preds = tf.reshape(
        tf.gather(y_pred, tf.range(start, start + bag_size), axis=0),
        shape=(-1,),
    )
    pseudo_labels = tf.stop_gradient(
        get_pseudo_labels(hard_labels, reg, bag_size, avg_true_label, preds)
    )
    loss += instance_loss(pseudo_labels, preds)
    start += bag_size
  return loss / batch_size
//  [python]
def erot_loss(
    hard_labels,
    reg,
    batch_size,
    instance_loss,
    x,
    y,
    y_pred,
):
  """EROT loss."""
  del x
  y_sg = tf.stop_gradient(y)
  loss = tf.zeros((), dtype=tf.float32)
  start = tf.zeros((), dtype=tf.int32)
  for _ in range(batch_size):
    bag_size = tf.cast(tf.gather(y_sg, start, axis=0)[0], dtype=tf.int32)
    avg_true_label = (
        tf.gather(y_sg, start + 1, axis=0)[0]
        / tf.gather(y_sg, start, axis=0)[0]
    )
    preds = tf.reshape(
        tf.gather(y_pred, tf.range(start, start + bag_size), axis=0),
        shape=(-1,),
    )
    pseudo_labels = tf.stop_gradient(
        get_pseudo_labels(hard_labels, reg, bag_size, avg_true_label, preds)
    )
    loss += instance_loss(pseudo_labels, preds)
    start += bag_size
  return loss / batch_size
//  [python]
def sink(
    a,
    b,
    # pylint: disable=invalid-name
    M,
    reg = 0.2,
    max_iter = 1000,
    error_thresh = 1e-9,
):
  """Sinkhorn algorithm in tensorflow."""
  a = tf.expand_dims(a, axis=1)
  b = tf.expand_dims(b, axis=1)
  u = tf.ones(shape=(tf.shape(a)[0], 1), dtype=tf.float32) / tf.to_float(
      tf.shape(a)[0]
  )
  v = tf.ones(shape=(tf.shape(b)[0], 1), dtype=tf.float32) / tf.to_float(
      tf.shape(b)[0]
  )
  kernel = tf.exp(-M / reg)
  kernel_p = (1.0 / a) * kernel
  counter = tf.constant(0)
  err = tf.constant(1.0)

  def cond(counter, u, v, err):
    del u, v
    return tf.logical_and(
        counter < max_iter,
        err > error_thresh,
    )

  def err_fn():
    return tf.pow(
        tf.norm(tf.reduce_sum(u * (kernel * tf.squeeze(v))) - b, ord=1), 2
    )

  def default_err():
    return err

  def loop_func(counter, u, v, err):
    del v, err
    v = tf.math.divide(b, tf.matmul(tf.transpose(kernel, (1, 0)), u))
    u = 1.0 / tf.matmul(kernel_p, v)
    err = tf.cond(tf.equal(counter % 10, 0), err_fn, default_err)
    counter = tf.add(counter, 1)
    return counter, u, v, err

  _, u, v, _ = tf.while_loop(cond, loop_func, loop_vars=[counter, u, v, err])
  return tf.reshape(u, (-1, 1)) * kernel * tf.reshape(v, (1, -1))
//  [python]
def sink(
    a,
    b,
    # pylint: disable=invalid-name
    M,
    reg = 0.2,
    max_iter = 1000,
    error_thresh = 1e-9,
):
  """Sinkhorn algorithm in tensorflow."""
  a = tf.expand_dims(a, axis=1)
  b = tf.expand_dims(b, axis=1)
  u = tf.ones(shape=(tf.shape(a)[0], 1), dtype=tf.float32) / tf.cast(
      tf.shape(a)[0], dtype=tf.float32
  )
  v = tf.ones(shape=(tf.shape(b)[0], 1), dtype=tf.float32) / tf.cast(
      tf.shape(b)[0], dtype=tf.float32
  )
  kernel = tf.exp(-M / reg)
  kernel_p = (1.0 / a) * kernel
  counter = tf.constant(0)
  err = tf.constant(1.0)

  def cond(counter, u, v, err):
    del u, v
    return tf.logical_and(
        counter < max_iter,
        err > error_thresh,
    )

  def err_fn():
    return tf.pow(
        tf.norm(tf.reduce_sum(u * (kernel * tf.squeeze(v))) - b, ord=1), 2
    )

  def default_err():
    return err

  def loop_func(counter, u, v, err):
    del v, err
    v = tf.math.divide(b, tf.matmul(tf.transpose(kernel, (1, 0)), u))
    u = 1.0 / tf.matmul(kernel_p, v)
    err = tf.cond(tf.equal(counter % 10, 0), err_fn, default_err)
    counter = tf.add(counter, 1)
    return counter, u, v, err

  _, u, v, _ = tf.while_loop(cond, loop_func, loop_vars=[counter, u, v, err])
  return tf.reshape(u, (-1, 1)) * kernel * tf.reshape(v, (1, -1))
//  [python]
def easy_llp_loss(
    batch_size,
    p,
    instance_loss,
    y,
    y_pred,
):
  """Easy-LLP loss."""
  p = tf.convert_to_tensor(p, dtype=tf.float32)
  start = 0
  loss = 0.0
  for _ in range(batch_size):
    bag_size = int(y[start, 0])
    if bag_size != 0:
      avg_true_label = tf.cast(y[start + 1, 0], tf.float32) / tf.cast(
          y[start, 0], tf.float32
      )
      idx = tf.random.uniform(shape=[], maxval=bag_size, dtype=tf.int32)
      r_pred = y_pred[start + idx, 0]
      loss += (y[start, 0] * (avg_true_label - p) + p) * instance_loss(
          tf.convert_to_tensor([1.0]), tf.convert_to_tensor([r_pred])
      ) + (y[start, 0] * (p - avg_true_label) + (1 - p)) * instance_loss(
          tf.convert_to_tensor([0.0]), tf.convert_to_tensor([r_pred])
      )
      start = start + bag_size
  return tf.convert_to_tensor(loss)
//  [python]
def call(self, inputs, training=None, masking=None):
    patches = inputs['patches']
    masked_indices = t.cast(inputs['masked_indices'], t.int32)
    unmasked_indices = t.cast(inputs['unmasked_indices'], t.int32)
    batch_size = tf.shape(patches)[0]
    num_h_patches = tf.shape(patches)[1]
    num_w_patches = tf.shape(patches)[2]
    num_patches = num_h_patches * num_w_patches
    num_masks = tf.shape(masked_indices)[1]
    patch_embeds = self.encoder.to_embed(patches)
    patch_embeds = self.encoder.add_position_embed(patch_embeds)
    patch_embeds = tf.reshape(
        patch_embeds,
        (batch_size, num_patches, -1))
    patch_embeds = self.encoder.insert_cls(patch_embeds)

    unmasked_indices = tf.concat(
        [tf.zeros((batch_size, 1), unmasked_indices.dtype),
         unmasked_indices + 1],
        axis=1)
    masked_indices = masked_indices + 1
    unmasked_patch_embeds = tf.gather(
        patch_embeds, unmasked_indices, batch_dims=1)
    encoded = self.encoder({'embeddings': unmasked_patch_embeds})
    encoded = self.linear(encoded)

    zeros = tf.zeros((batch_size, num_patches + 1, 512))

    unmasked_embed = tf.tensor_scatter_nd_add(
        zeros,
        tf.stack([
            tf.tile(
                tf.expand_dims(tf.range(batch_size), axis=1),
                [1, num_patches + 1 - num_masks]), unmasked_indices
        ],
                 axis=-1),
        encoded)
    mask_embeds = tf.tile(self.mask, [batch_size, num_masks, 1])
    full_embed = tf.tensor_scatter_nd_add(
        unmasked_embed,
        tf.stack([
            tf.tile(
                tf.expand_dims(tf.range(batch_size), axis=1),
                [1, num_masks]), masked_indices
        ],
                 axis=-1),
        mask_embeds)
    full_embed = self.add_position_embed(
        full_embed, num_h_patches, num_w_patches)

    decoded = self.decoder(full_embed)
    pred_pixel_values = self.to_pixels(
        tf.gather(decoded, masked_indices, batch_dims=1))
    return pred_pixel_values
//  [python]
def call(self, inputs, training=None, masking=None):
    patches = inputs['patches']
    masked_indices = tf.cast(inputs['masked_indices'], tf.int32)
    unmasked_indices = tf.cast(inputs['unmasked_indices'], tf.int32)
    batch_size = tf.shape(patches)[0]
    num_h_patches = tf.shape(patches)[1]
    num_w_patches = tf.shape(patches)[2]
    num_patches = num_h_patches * num_w_patches
    num_masks = tf.shape(masked_indices)[1]
    patch_embeds = self.encoder.to_embed(patches)
    patch_embeds = self.encoder.add_position_embed(patch_embeds)
    patch_embeds = tf.reshape(
        patch_embeds,
        (batch_size, num_patches, -1))
    patch_embeds = self.encoder.insert_cls(patch_embeds)

    unmasked_indices = tf.concat(
        [tf.zeros((batch_size, 1), unmasked_indices.dtype),
         unmasked_indices + 1],
        axis=1)
    masked_indices = masked_indices + 1
    unmasked_patch_embeds = tf.gather(
        patch_embeds, unmasked_indices, batch_dims=1)
    encoded = self.encoder({'embeddings': unmasked_patch_embeds})
    encoded = self.linear(encoded)

    zeros = tf.zeros((batch_size, num_patches + 1, 512))

    unmasked_embed = tf.tensor_scatter_nd_add(
        zeros,
        tf.stack([
            tf.tile(
                tf.expand_dims(tf.range(batch_size), axis=1),
                [1, num_patches + 1 - num_masks]), unmasked_indices
        ],
                 axis=-1),
        encoded)
    mask_embeds = tf.tile(self.mask, [batch_size, num_masks, 1])
    full_embed = tf.tensor_scatter_nd_add(
        unmasked_embed,
        tf.stack([
            tf.tile(
                tf.expand_dims(tf.range(batch_size), axis=1),
                [1, num_masks]), masked_indices
        ],
                 axis=-1),
        mask_embeds)
    full_embed = self.add_position_embed(
        full_embed, num_h_patches, num_w_patches)

    decoded = self.decoder(full_embed)
    pred_pixel_values = self.to_pixels(
        tf.gather(decoded, masked_indices, batch_dims=1))
    return pred_pixel_values
//  [python]
def ot_llp_loss(
    batch_size,
    instance_loss,
    y,
    y_pred,
):
  """OT without ER loss."""
  start = 0
  loss = 0.0
  for _ in range(batch_size):
    bag_size = int(y[start, 0])
    sum_true_label = int(y[start + 1, 0])
    preds = y_pred[start : start + bag_size, 0]
    if bag_size != 0:
      if sum_true_label == 0:
        pseudo_labels = tf.zeros(tf.shape(preds), dtype=tf.float32)
      else:
        values, _ = tf.math.top_k(preds, k=sum_true_label)
        pseudo_labels = tf.to_float(preds >= values[-1])
      loss += instance_loss(pseudo_labels, preds)
      start = start + bag_size
  return tf.convert_to_tensor(loss)
//  [python]
def ot_llp_loss(
    batch_size,
    instance_loss,
    y,
    y_pred,
):
  """OT without ER loss."""
  start = 0
  loss = 0.0
  for _ in range(batch_size):
    bag_size = int(y[start, 0])
    sum_true_label = int(y[start + 1, 0])
    preds = y_pred[start : start + bag_size, 0]
    if bag_size != 0:
      if sum_true_label == 0:
        pseudo_labels = tf.zeros(tf.shape(preds), dtype=tf.float32)
      else:
        values, _ = tf.math.top_k(preds, k=sum_true_label)
        pseudo_labels = tf.cast(preds >= values[-1], tf.float32)
      loss += instance_loss(pseudo_labels, preds)
      start = start + bag_size
  return tf.convert_to_tensor(loss)
//  [python]
def dllp_loss_graph_regression(
    batch_size,
    instance_loss,
    x_categ,
    x_numer,
    y,
    y_pred,
):
  """DLLP loss for graph execution."""
  del x_categ, x_numer
  y_sg = tf.stop_gradient(y)
  loss = tf.zeros((), dtype=tf.float32)
  start = tf.zeros((), dtype=tf.int32)
  for _ in range(batch_size):
    bag_size = tf.to_int32(tf.gather(y_sg, start, axis=0)[0])
    avg_true_label = tf.gather(y_sg, start + 1, axis=0)[0]
    avg_pred_label = tf.reduce_mean(
        tf.gather(y_pred, tf.range(start, start + bag_size), axis=0)
    )
    loss += instance_loss(
        tf.convert_to_tensor([avg_true_label]),
        tf.convert_to_tensor([avg_pred_label]),
    )
    start += bag_size
  return loss / batch_size
//  [python]
def dllp_loss_graph_regression(
    batch_size,
    instance_loss,
    x_categ,
    x_numer,
    y,
    y_pred,
):
  """DLLP loss for graph execution."""
  del x_categ, x_numer
  y_sg = tf.stop_gradient(y)
  loss = tf.zeros((), dtype=tf.float32)
  start = tf.zeros((), dtype=tf.int32)
  for _ in range(batch_size):
    bag_size = tf.cast(tf.gather(y_sg, start, axis=0)[0], dtype=tf.int32)
    avg_true_label = tf.gather(y_sg, start + 1, axis=0)[0]
    avg_pred_label = tf.reduce_mean(
        tf.gather(y_pred, tf.range(start, start + bag_size), axis=0)
    )
    loss += instance_loss(
        tf.convert_to_tensor([avg_true_label]),
        tf.convert_to_tensor([avg_pred_label]),
    )
    start += bag_size
  return loss / batch_size
//  [python]
def frame_preprocess(frame):

    def _frame_preprocess(frame):
        """Breakout向けの切り取りであることに注意
        """
        image = tf.to_float(tf.convert_to_tensor(frame))
        image_gray = tf.image.rgb_to_grayscale(image)
        image_crop = tf.image.crop_to_bounding_box(image_gray, 34, 0, 160, 160)
        image_resize = tf.image.resize(image_crop, [84, 84])
        image_scaled = tf.divide(image_resize, 255)
        return image_scaled

    frame = _frame_preprocess(frame).numpy()[:, :, 0]

    return frame
//  [python]
def frame_preprocess(frame):

    def _frame_preprocess(frame):
        """Breakout向けの切り取りであることに注意
        """
        image = tf.cast(tf.convert_to_tensor(frame), tf.float32)
        image_gray = tf.image.rgb_to_grayscale(image)
        image_crop = tf.image.crop_to_bounding_box(image_gray, 34, 0, 160, 160)
        image_resize = tf.image.resize(image_crop, [84, 84])
        image_scaled = tf.divide(image_resize, 255)
        return image_scaled

    frame = _frame_preprocess(frame).numpy()[:, :, 0]

    return frame
//  [python]
def set_with_mask(x, x_other, mask):
    """Helper function which returns a tensor similar to x with all the values
    of x replaced by x_other where the mask evaluates to true.
    """
    mask = t.cast(mask, x.dtype)
    ones = tf.ones_like(mask, dtype=x.dtype)
    return x_other * mask + x * (ones - mask)
//  [python]
def set_with_mask(x, x_other, mask):
    """Helper function which returns a tensor similar to x with all the values
    of x replaced by x_other where the mask evaluates to true.
    """
    mask = tf.cast(mask, x.dtype)
    ones = tf.ones_like(mask, dtype=x.dtype)
    return x_other * mask + x * (ones - mask)
//  [python]
def genbags_loss_regression(
    mean,
    covariance_mat,
    batch_size,
    block_size,
    num_gen_bags_per_block,
    x_categ,
    x_numer,
    y,
    y_pred,
):
  """Genbags loss."""
  del x_categ, x_numer
  loss = tf.zeros((), dtype=tf.float32)
  start = tf.zeros((), dtype=tf.int32)
  for _ in range(batch_size // block_size):
    wts = tfp.distributions.MultivariateNormalFullCovariance(
        loc=mean, covariance_matrix=covariance_mat
    ).sample(num_gen_bags_per_block)
    list_of_diffs = []
    for _ in range(block_size):
      bag_size = tf.to_int32(tf.gather(y, start, axis=0)[0])
      avg_true_label = tf.gather(y, start + 1, axis=0)[0]
      avg_pred_label = tf.reduce_mean(
          tf.gather(y_pred, tf.range(start, start + bag_size), axis=0)
      )
      list_of_diffs.append(avg_pred_label - avg_true_label)
      start = start + bag_size
    loss = loss + tf.square(
        tf.norm(
            tf.linalg.matvec(wts, tf.convert_to_tensor(list_of_diffs)), ord=2
        )
    )
  return loss
//  [python]
def genbags_loss_regression(
    mean,
    covariance_mat,
    batch_size,
    block_size,
    num_gen_bags_per_block,
    x_categ,
    x_numer,
    y,
    y_pred,
):
  """Genbags loss."""
  del x_categ, x_numer
  loss = tf.zeros((), dtype=tf.float32)
  start = tf.zeros((), dtype=tf.int32)
  for _ in range(batch_size // block_size):
    wts = tfp.distributions.MultivariateNormalFullCovariance(
        loc=mean, covariance_matrix=covariance_mat
    ).sample(num_gen_bags_per_block)
    list_of_diffs = []
    for _ in range(block_size):
      bag_size = tf.cast(tf.gather(y, start, axis=0)[0], dtype=tf.int32)
      avg_true_label = tf.gather(y, start + 1, axis=0)[0]
      avg_pred_label = tf.reduce_mean(
          tf.gather(y_pred, tf.range(start, start + bag_size), axis=0)
      )
      list_of_diffs.append(avg_pred_label - avg_true_label)
      start = start + bag_size
    loss = loss + tf.square(
        tf.norm(
            tf.linalg.matvec(wts, tf.convert_to_tensor(list_of_diffs)), ord=2
        )
    )
  return loss
//  [python]
def get_or_guess_labels(model_fn, x, y=None, targeted=False):
    """
    Helper function to get the label to use in generating an
    adversarial example for x.
    If 'y' is not None, then use these labels.
    If 'targeted' is True, then assume it's a targeted attack
    and y must be set.
    Otherwise, use the model's prediction as the label and perform an
    untargeted attack
    :param model_fn: a callable that takes an input tensor and returns the model logits.
    :param x: input tensor.
    """
    if targeted is True and y is None:
        raise ValueError("Must provide y for a targeted attack!")

    preds = model_fn(x)
    nb_classes = preds.shape[-1]

    # labels set by the user
    if y is not None:
        # inefficient when y is a tensor, but this function only get called once
        y = np.asarray(y)

        if len(y.shape) == 1:
            # the user provided categorical encoding
            y = tf.one_hot(y, nb_classes)

        y = tf.to_float(y)
        return y, nb_classes

    # must be an untargeted attack
    labels = tf.to_float(
        tf.equal(tf.reduce_max(preds, axis=1, keepdims=True), preds)
    )

    return labels, nb_classes
//  [python]
def get_or_guess_labels(model_fn, x, y=None, targeted=False):
    """
    Helper function to get the label to use in generating an
    adversarial example for x.
    If 'y' is not None, then use these labels.
    If 'targeted' is True, then assume it's a targeted attack
    and y must be set.
    Otherwise, use the model's prediction as the label and perform an
    untargeted attack
    :param model_fn: a callable that takes an input tensor and returns the model logits.
    :param x: input tensor.
    """
    if targeted is True and y is None:
        raise ValueError("Must provide y for a targeted attack!")

    preds = model_fn(x)
    nb_classes = preds.shape[-1]

    # labels set by the user
    if y is not None:
        # inefficient when y is a tensor, but this function only get called once
        y = np.asarray(y)

        if len(y.shape) == 1:
            # the user provided categorical encoding
            y = tf.one_hot(y, nb_classes)

        y = tf.cast(y, x.dtype)
        return y, nb_classes

    # must be an untargeted attack
    labels = tf.cast(
        tf.equal(tf.reduce_max(preds, axis=1, keepdims=True), preds), x.dtype
    )

    return labels, nb_classes
//  [python]
def get_loss(self, logits, labels, activation, pre_cal_weights):
        # calculate the weighted cross entropy according to the inverse frequency
        class_weights = tf.convert_to_tensor(pre_cal_weights, dtype=tf.float32)
        one_hot_labels = tf.one_hot(labels, depth=self.config.num_classes)


        self.ssdr_one_hot_labels = one_hot_labels
        self.ssdr_logits = logits


        unweighted_losses = tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=one_hot_labels)

        weights = tf.reduce_sum(class_weights * one_hot_labels, axis=1)

        weighted_losses = unweighted_losses * weights
        weighted_losses_acti = weighted_losses * tf.to_float(activation)
        output_loss = tf.reduce_mean(weighted_losses_acti)
        return output_loss
//  [python]
def get_loss(self, logits, labels, activation, pre_cal_weights):
        # calculate the weighted cross entropy according to the inverse frequency
        class_weights = tf.convert_to_tensor(pre_cal_weights, dtype=tf.float32)
        one_hot_labels = tf.one_hot(labels, depth=self.config.num_classes)


        self.ssdr_one_hot_labels = one_hot_labels
        self.ssdr_logits = logits


        unweighted_losses = tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=one_hot_labels)

        weights = tf.reduce_sum(class_weights * one_hot_labels, axis=1)

        weighted_losses = unweighted_losses * weights
        weighted_losses_acti = weighted_losses * tf.cast(activation, dtype=tf.float32)
        output_loss = tf.reduce_mean(weighted_losses_acti)
        return output_loss
//  [python]
def _cross_suppression(boxes, box_slice, iou_threshold, inner_idx):
  batch_size = tf.shape(boxes)[0]
  new_slice = tf.slice(boxes, [0, inner_idx * NMS_TILE_SIZE, 0],
                       [batch_size, NMS_TILE_SIZE, 4])
  iou = box_utils.bbox_overlap(new_slice, box_slice)
  ret_slice = tf.expand_dims(
      tf.to_float(tf.reduce_all(iou < iou_threshold, [1])),
      2) * box_slice
  return boxes, ret_slice, iou_threshold, inner_idx + 1
//  [python]
def _cross_suppression(boxes, box_slice, iou_threshold, inner_idx):
  batch_size = tf.shape(boxes)[0]
  new_slice = tf.slice(boxes, [0, inner_idx * NMS_TILE_SIZE, 0],
                       [batch_size, NMS_TILE_SIZE, 4])
  iou = box_utils.bbox_overlap(new_slice, box_slice)
  ret_slice = tf.expand_dims(
      tf.cast(tf.reduce_all(iou < iou_threshold, [1]), box_slice.dtype),
      2) * box_slice
  return boxes, ret_slice, iou_threshold, inner_idx + 1
//  [python]
def _self_suppression(iou, _, iou_sum):
  batch_size = tf.shape(iou)[0]
  can_suppress_others = tf.to_float(
      tf.reshape(tf.reduce_max(iou, 1) <= 0.5, [batch_size, -1, 1]))
  iou_suppressed = tf.reshape(
      tf.to_float(tf.reduce_max(can_suppress_others * iou, 1) <= 0.5),
      [batch_size, -1, 1]) * iou
  iou_sum_new = tf.reduce_sum(iou_suppressed, [1, 2])
  return [
      iou_suppressed,
      tf.reduce_any(iou_sum - iou_sum_new > 0.5), iou_sum_new
  ]
//  [python]
def _self_suppression(iou, _, iou_sum):
  batch_size = tf.shape(iou)[0]
  can_suppress_others = tf.cast(
      tf.reshape(tf.reduce_max(iou, 1) <= 0.5, [batch_size, -1, 1]), iou.dtype)
  iou_suppressed = tf.reshape(
      tf.cast(tf.reduce_max(can_suppress_others * iou, 1) <= 0.5, iou.dtype),
      [batch_size, -1, 1]) * iou
  iou_sum_new = tf.reduce_sum(iou_suppressed, [1, 2])
  return [
      iou_suppressed,
      tf.reduce_any(iou_sum - iou_sum_new > 0.5), iou_sum_new
  ]
//  [python]
def create_padding_mask(seq):
    seq = tf.to_float(tf.math.equal(seq, 0))
  
    # add extra dimensions to add the padding
    # to the attention logits.
    return seq[:, tf.newaxis, tf.newaxis, :]
//  [python]
def create_padding_mask(seq):
    seq = tf.cast(tf.math.equal(seq, 0), tf.float32)
  
    # add extra dimensions to add the padding
    # to the attention logits.
    return seq[:, tf.newaxis, tf.newaxis, :]
//  [python]
def positional_encoding(position, d_model):
    angle_rads = get_angles(np.arange(position)[:, np.newaxis], np.arange(d_model)[np.newaxis, :],d_model)
  
    # apply sin to even indices in the array; 2i
    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])
  
    # apply cos to odd indices in the array; 2i+1
    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])
    
    pos_encoding = angle_rads[np.newaxis, ...]
    
    return t.cast(pos_encoding, dtype=t.float32)
//  [python]
def positional_encoding(position, d_model):
    angle_rads = get_angles(np.arange(position)[:, np.newaxis], np.arange(d_model)[np.newaxis, :],d_model)
  
    # apply sin to even indices in the array; 2i
    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])
  
    # apply cos to odd indices in the array; 2i+1
    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])
    
    pos_encoding = angle_rads[np.newaxis, ...]
    
    return tf.cast(pos_encoding, dtype=tf.float32)