jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def compute_arsm_grads(
      self, input_batch, encoder_logits, dirichlet_samples,
      partial_integration=False):
    batch_size, num_variables, num_classes = encoder_logits.shape
    # [batch_size, num_classes, num_classes, num_variables]
    swapped_matrix = self.get_arsm_swapped_sample(
        encoder_logits, dirichlet_samples)
    swapped_matrix = tf.reshape(
        swapped_matrix,
        [batch_size * num_classes * num_classes, num_variables])
    # input_batch is of the size [batch_size, observation_dims]
    tiled_inputs = tf.tile(input_batch[:, None, None, :],
                           [1, num_classes, num_classes, 1])
    tiled_inputs = tf.reshape(
        tiled_inputs,
        [batch_size * num_classes * num_classes, -1])
    loss_matrix = self.get_elbo(tiled_inputs, swapped_matrix)
    loss_matrix = tf.reshape(loss_matrix,
                             [batch_size, num_classes, num_classes])

    tilde_loss = (loss_matrix
                  - tf.math.reduce_mean(loss_matrix, axis=-1, keepdims=True))
    tilde_dirichlet = 1./num_classes - dirichlet_samples
    if partial_integration:
      swapped_matrix = tf.reshape(
          swapped_matrix,
          [batch_size, num_classes, num_classes, num_variables])
      tilde_dirichlet *= tf.transpose(
          tf.cast(tf.logical_not(tf.reduce_all(
              swapped_matrix == swapped_matrix[:, :, 0, :][:, :, None, :],
              axis=2)), dtype=tf.float32),
          [0, 2, 1])

    layer_grad = tf.matmul(tilde_dirichlet, tilde_loss)
    return layer_grad
//  [python]
def compute_ars_grads(
      self, input_batch, encoder_logits, dirichlet_samples,
      partial_integration=False):
    batch_size, num_variables, num_classes = encoder_logits.shape
    if partial_integration:
      swapped_vector, reference_samples = self.get_arsplus_swapped_vector(
          encoder_logits, dirichlet_samples)
    else:
      swapped_vector, reference_samples = self.get_ars_swapped_vector(
          encoder_logits, dirichlet_samples)
    swapped_vector = tf.reshape(
        swapped_vector,
        [batch_size * num_classes, num_variables])
    # input_batch is of the size [batch_size, observation_dims]
    tiled_inputs = tf.tile(input_batch[:, None, :],
                           [1, num_classes, 1])
    tiled_inputs = tf.reshape(
        tiled_inputs,
        [batch_size * num_classes, -1])
    loss_vector = self.get_elbo(tiled_inputs, swapped_vector)
    loss_vector = tf.reshape(loss_vector,
                             [batch_size, num_classes])
    swapped_vector = tf.reshape(swapped_vector,
                                [batch_size, num_classes, num_variables])

    tilde_loss = (loss_vector
                  - tf.math.reduce_mean(loss_vector, axis=-1, keepdims=True))
    tilde_dirichlet = 1. - num_classes * reference_samples
    if partial_integration:
      tilde_dirichlet *= tf.cast(tf.logical_not(tf.reduce_all(
          swapped_vector == swapped_vector[:, 0, :][:, None, :],
          axis=1)), dtype=tf.float32)
    layer_grad = tilde_loss[:, None, :] * tilde_dirichlet[:, :, None]
    return layer_grad
//  [python]
def build(self, input_shape: tf.TensorShape) -> None:
        points = list(itertools.product(range(self.resolution), range(self.resolution)))
        num_points = len(points)
        attention_offsets = {}

        idxs = []

        for point_1 in points:
            for point_2 in points:
                offset = (abs(point_1[0] - point_2[0]), abs(point_1[1] - point_2[1]))
                if offset not in attention_offsets:
                    attention_offsets[offset] = len(attention_offsets)
                idxs.append(attention_offsets[offset])

        self.attention_biases = self.add_weight(
            shape=(self.num_heads, len(attention_offsets)),
            initializer=keras.initializers.zeros(),
            trainable=True,
            name="attention_biases",
        )
        self.attention_bias_idxs = self.add_weight(
            shape=(num_points, num_points),
            trainable=False,
            dtype=tf.int32,
            name="attention_bias_idxs",
        )

        self.attention_bias_idxs.assign(tf.reshape(tf.to_int32(idxs), (num_points, num_points)))

        if self.built:
            return
        self.built = True
        if getattr(self, "qkv", None) is not None:
            with tf.name_scope(self.qkv.name):
                self.qkv.build([None, None, self.dim])
        if getattr(self, "projection", None) is not None:
            with tf.name_scope(self.projection.name):
                self.projection.build([None, None, self.total_expanded_key_dim])
//  [python]
def build(self, input_shape: tf.TensorShape) -> None:
        points = list(itertools.product(range(self.resolution), range(self.resolution)))
        num_points = len(points)
        attention_offsets = {}

        idxs = []

        for point_1 in points:
            for point_2 in points:
                offset = (abs(point_1[0] - point_2[0]), abs(point_1[1] - point_2[1]))
                if offset not in attention_offsets:
                    attention_offsets[offset] = len(attention_offsets)
                idxs.append(attention_offsets[offset])

        self.attention_biases = self.add_weight(
            shape=(self.num_heads, len(attention_offsets)),
            initializer=keras.initializers.zeros(),
            trainable=True,
            name="attention_biases",
        )
        self.attention_bias_idxs = self.add_weight(
            shape=(num_points, num_points),
            trainable=False,
            dtype=tf.int32,
            name="attention_bias_idxs",
        )

        self.attention_bias_idxs.assign(tf.reshape(tf.cast(idxs, dtype=tf.int32), (num_points, num_points)))

        if self.built:
            return
        self.built = True
        if getattr(self, "qkv", None) is not None:
            with tf.name_scope(self.qkv.name):
                self.qkv.build([None, None, self.dim])
        if getattr(self, "projection", None) is not None:
            with tf.name_scope(self.projection.name):
                self.projection.build([None, None, self.total_expanded_key_dim])
//  [python]
def call(
        self, hidden_states: tf.Tensor, output_attentions: bool = False, training: bool = False
    ) -> Tuple[tf.Tensor]:
        batch_size, sequence_length, *_ = shape_list(hidden_states)
        qkv = self.qkv(inputs=hidden_states)

        query_layer, key_layer, value_layer = tf.split(
            tf.reshape(tensor=qkv, shape=(batch_size, sequence_length, self.num_heads, -1)),
            num_or_size_splits=[self.key_dim, self.key_dim, self.expanded_key_dim],
            axis=3,
        )

        query_layer = tf.transpose(query_layer, perm=[0, 2, 1, 3])
        key_layer = tf.transpose(key_layer, perm=[0, 2, 1, 3])
        value_layer = tf.transpose(value_layer, perm=[0, 2, 1, 3])

        attention_probs = tf.matmul(query_layer, tf.transpose(key_layer, perm=[0, 1, 3, 2]))
        scale = tf.to_float(self.scale)
        attention_probs = tf.multiply(attention_probs, scale)

        attention_biases = tf.gather(params=self.attention_biases, indices=self.attention_bias_idxs, axis=1)
        attention_probs = attention_probs + attention_biases
        attention_probs = stable_softmax(logits=attention_probs, axis=-1)

        context_layer = tf.matmul(attention_probs, value_layer)
        context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])

        context_layer = tf.reshape(
            tensor=context_layer, shape=(batch_size, sequence_length, self.total_expanded_key_dim)
        )
        context_layer = self.projection(context_layer)

        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)

        return outputs
//  [python]
def call(
        self, hidden_states: tf.Tensor, output_attentions: bool = False, training: bool = False
    ) -> Tuple[tf.Tensor]:
        batch_size, sequence_length, *_ = shape_list(hidden_states)
        qkv = self.qkv(inputs=hidden_states)

        query_layer, key_layer, value_layer = tf.split(
            tf.reshape(tensor=qkv, shape=(batch_size, sequence_length, self.num_heads, -1)),
            num_or_size_splits=[self.key_dim, self.key_dim, self.expanded_key_dim],
            axis=3,
        )

        query_layer = tf.transpose(query_layer, perm=[0, 2, 1, 3])
        key_layer = tf.transpose(key_layer, perm=[0, 2, 1, 3])
        value_layer = tf.transpose(value_layer, perm=[0, 2, 1, 3])

        attention_probs = tf.matmul(query_layer, tf.transpose(key_layer, perm=[0, 1, 3, 2]))
        scale = tf.cast(self.scale, dtype=attention_probs.dtype)
        attention_probs = tf.multiply(attention_probs, scale)

        attention_biases = tf.gather(params=self.attention_biases, indices=self.attention_bias_idxs, axis=1)
        attention_probs = attention_probs + attention_biases
        attention_probs = stable_softmax(logits=attention_probs, axis=-1)

        context_layer = tf.matmul(attention_probs, value_layer)
        context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])

        context_layer = tf.reshape(
            tensor=context_layer, shape=(batch_size, sequence_length, self.total_expanded_key_dim)
        )
        context_layer = self.projection(context_layer)

        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)

        return outputs
//  [python]
def meshgrid(batch, height, width, is_homogeneous=True):
    """Construct a 2D meshgrid.

    Args:
      batch: batch size
      height: height of the grid
      width: width of the grid
      is_homogeneous: whether to return in homogeneous coordinates
    Returns:
      x,y grid coordinates [batch, 2 (3 if homogeneous), height, width]
    """
    x_t = tf.matmul(tf.ones(shape=tf.stack([height, 1])),
                    tf.transpose(tf.expand_dims(
                        tf.linspace(-1.0, 1.0, width), 1), [1, 0]))
    y_t = tf.matmul(tf.expand_dims(tf.linspace(-1.0, 1.0, height), 1),
                    tf.ones(shape=tf.stack([1, width])))
    x_t = (x_t + 1.0) * 0.5 * tf.to_float(width - 1)
    y_t = (y_t + 1.0) * 0.5 * tf.to_float(height - 1)
    if is_homogeneous:
        ones = tf.ones_like(x_t)
        coords = tf.stack([x_t, y_t, ones], axis=0)
    else:
        coords = tf.stack([x_t, y_t], axis=0)
    coords = tf.tile(tf.expand_dims(coords, 0), [batch, 1, 1, 1])
    return coords
//  [python]
def meshgrid(batch, height, width, is_homogeneous=True):
    """Construct a 2D meshgrid.

    Args:
      batch: batch size
      height: height of the grid
      width: width of the grid
      is_homogeneous: whether to return in homogeneous coordinates
    Returns:
      x,y grid coordinates [batch, 2 (3 if homogeneous), height, width]
    """
    x_t = tf.matmul(tf.ones(shape=tf.stack([height, 1])),
                    tf.transpose(tf.expand_dims(
                        tf.linspace(-1.0, 1.0, width), 1), [1, 0]))
    y_t = tf.matmul(tf.expand_dims(tf.linspace(-1.0, 1.0, height), 1),
                    tf.ones(shape=tf.stack([1, width])))
    x_t = (x_t + 1.0) * 0.5 * tf.cast(width - 1, tf.float32)
    y_t = (y_t + 1.0) * 0.5 * tf.cast(height - 1, tf.float32)
    if is_homogeneous:
        ones = tf.ones_like(x_t)
        coords = tf.stack([x_t, y_t, ones], axis=0)
    else:
        coords = tf.stack([x_t, y_t], axis=0)
    coords = tf.tile(tf.expand_dims(coords, 0), [batch, 1, 1, 1])
    return coords
//  [python]
def compute_logit_length(self, samples):
        input_length = tf.to_float(samples["input_length"])
        logit_length = tf.ceil(input_length / 2)
        logit_length = tf.ceil(logit_length / 2)
        logit_length = tf.to_int32(logit_length)
        return logit_length
//  [python]
def compute_logit_length(self, samples):
        input_length = tf.cast(samples["input_length"], tf.float32)
        logit_length = tf.math.ceil(input_length / 2)
        logit_length = tf.math.ceil(logit_length / 2)
        logit_length = tf.cast(logit_length, tf.int32)
        return logit_length
//  [python]
def _rescale(self, g):
        """
        Scale the gradients for plotting

        :param Tensor g: tensor to scale
        """
        try:
            output = g * tf.cast(self._loss_scale_value, g.dtype)
        except Exception as e:
            tf.compat.v1.logging.debug(e)
            output = g
        return output
//  [python]
def generate_mpc_mask(self, input_data):
        """ generate mask for pretraining

        Args:
            acoustic features: i.e F-bank
        Return::

            mask tensor
        """
        dtype = input_data.dtype
        chunk_size = self.hparams.chunk_size * self.downsample_scale
        batch, seq_len, dim, num_channels = input_data.get_shape().as_list()
        if (1 - self.hparams.keep_probability) * seq_len <= chunk_size:
            chunk_size = self.downsample_scale
        num_chunk = tf.to_int32(
            tf.ceil(
                tf.div(tf.to_float(seq_len), tf.to_float(chunk_size))
            )
        )

        # generate mask with shape [batch, num_chunk]: 1.0 for keep, 0.0 for masked
        ones = tf.ones([batch, num_chunk], dtype=dtype)
        zeros = tf.zeros([batch, num_chunk], dtype=dtype)
        probability = ones * self.hparams.keep_probability
        random = self.randomizer([batch, num_chunk], dtype=dtype)
        mask = tf.where(tf.less(random, probability), ones, zeros)

        # change the mask for [batch, seq_len]
        mask = tf.tile(mask[:, :, tf.newaxis], [1, 1, chunk_size])
        mask = tf.reshape(mask, [batch, -1])[:, :seq_len]

        # change the mask into masks with the shape [batch, seq_len, dim, num_channels]
        masks = tf.tile(mask[:, :, tf.newaxis, tf.newaxis], [1, 1, dim, num_channels])
        return masks
//  [python]
def generate_mpc_mask(self, input_data):
        """ generate mask for pretraining

        Args:
            acoustic features: i.e F-bank
        Return::

            mask tensor
        """
        dtype = input_data.dtype
        chunk_size = self.hparams.chunk_size * self.downsample_scale
        batch, seq_len, dim, num_channels = input_data.get_shape().as_list()
        if (1 - self.hparams.keep_probability) * seq_len <= chunk_size:
            chunk_size = self.downsample_scale
        num_chunk = tf.cast(
            tf.math.ceil(
                tf.divide(tf.cast(seq_len, tf.float32), tf.cast(chunk_size, tf.float32))
            ),
            tf.int32,
        )

        # generate mask with shape [batch, num_chunk]: 1.0 for keep, 0.0 for masked
        ones = tf.ones([batch, num_chunk], dtype=dtype)
        zeros = tf.zeros([batch, num_chunk], dtype=dtype)
        probability = ones * self.hparams.keep_probability
        random = self.randomizer([batch, num_chunk], dtype=dtype)
        mask = tf.where(tf.less(random, probability), ones, zeros)

        # change the mask for [batch, seq_len]
        mask = tf.tile(mask[:, :, tf.newaxis], [1, 1, chunk_size])
        mask = tf.reshape(mask, [batch, -1])[:, :seq_len]

        # change the mask into masks with the shape [batch, seq_len, dim, num_channels]
        masks = tf.tile(mask[:, :, tf.newaxis, tf.newaxis], [1, 1, dim, num_channels])
        return masks
//  [python]
def build_model(self):
        with tf.variable_scope('AttRec'):
            # Embedding
            self.item_emb = tf.get_variable("item_emb", [self.item_count, self.emb_size],initializer=self.u_init) #[N,e]

            self.item_rep_emb = tf.get_variable("item_rep_emb", [self.item_count, self.emb_size],initializer=self.u_init) #[N,e]
            self.user_rep_emb = tf.get_variable("user_rep_emb", [self.user_count, self.emb_size],initializer=self.u_init) #[N,e]

            value = tf.nn.embedding_lookup(self.item_emb, self.hist_seq,max_norm=1)  # [B,T,e]

            value = self.mask_seq(value,self.sl)
            #Add TimeSignals
            ts = self.make_time_signal(self.emb_size, self.msl) #[msl, e]
            ts = tf.tile(tf.expand_dims(ts, [0]), [tf.shape(self.sl)[0], 1, 1])

            query = tf.add(value, ts) #[B,T,e]
            key = query

            #Self Attention
            output = self.attention_module(query,key,value,self.emb_size)

            # Mean
            div_num = tf.to_float(tf.tile(tf.expand_dims(self.sl, 1), [1, self.emb_size]))
            m = tf.div(tf.reduce_sum(output, axis=1), div_num) #[B,e]
            
            #Look up embedding for targets and negative samples
            u = tf.nn.embedding_lookup(self.user_rep_emb, self.u_p) #[B,e]
            u = tf.nn.l2_normalize(u,-1)
            u = tf.clip_by_norm(u, 1, -1)

            pos_v = tf.nn.embedding_lookup(self.item_rep_emb, self.next_p) #[B,pos,e]
            pos_v = tf.clip_by_norm(pos_v,1,-1)

            neg_v = tf.nn.embedding_lookup(self.item_rep_emb, self.neg_p) #[B,neg,e]
            neg_v = tf.clip_by_norm(neg_v, 1, -1)

            pos_x = tf.nn.embedding_lookup(self.item_emb, self.next_p) #[B,pos,e]
            pos_x = tf.clip_by_norm(pos_x, 1, -1)


            neg_x = tf.nn.embedding_lookup(self.item_emb, self.neg_p)#[B,neg,e]
            neg_x = tf.clip_by_norm(neg_x, 1, -1)

            self.pos_y = pos_y = self.pos_object_function(u, pos_v, m, pos_x, self.w) #[B,tsl]
            self.neg_y = neg_y = self.neg_object_function(u, neg_v, m, neg_x, self.w) #[B,nsl]

            #margin based hinge loss
            self.loss = self.loss_function(self.gamma, pos_y, neg_y)
            self.next_items = m
//  [python]
def build_model(self):
        with tf.variable_scope('AttRec'):
            # Embedding
            self.item_emb = tf.get_variable("item_emb", [self.item_count, self.emb_size],initializer=self.u_init) #[N,e]

            self.item_rep_emb = tf.get_variable("item_rep_emb", [self.item_count, self.emb_size],initializer=self.u_init) #[N,e]
            self.user_rep_emb = tf.get_variable("user_rep_emb", [self.user_count, self.emb_size],initializer=self.u_init) #[N,e]

            value = tf.nn.embedding_lookup(self.item_emb, self.hist_seq,max_norm=1)  # [B,T,e]

            value = self.mask_seq(value,self.sl)
            #Add TimeSignals
            ts = self.make_time_signal(self.emb_size, self.msl) #[msl, e]
            ts = tf.tile(tf.expand_dims(ts, [0]), [tf.shape(self.sl)[0], 1, 1])

            query = tf.add(value, ts) #[B,T,e]
            key = query

            #Self Attention
            output = self.attention_module(query,key,value,self.emb_size)

            # Mean
            div_num = tf.cast(tf.tile(tf.expand_dims(self.sl, 1), [1, self.emb_size]), tf.float32)
            m = tf.div(tf.reduce_sum(output, axis=1), div_num) #[B,e]
            
            #Look up embedding for targets and negative samples
            u = tf.nn.embedding_lookup(self.user_rep_emb, self.u_p) #[B,e]
            u = tf.nn.l2_normalize(u,-1)
            u = tf.clip_by_norm(u, 1, -1)

            pos_v = tf.nn.embedding_lookup(self.item_rep_emb, self.next_p) #[B,pos,e]
            pos_v = tf.clip_by_norm(pos_v,1,-1)

            neg_v = tf.nn.embedding_lookup(self.item_rep_emb, self.neg_p) #[B,neg,e]
            neg_v = tf.clip_by_norm(neg_v, 1, -1)

            pos_x = tf.nn.embedding_lookup(self.item_emb, self.next_p) #[B,pos,e]
            pos_x = tf.clip_by_norm(pos_x, 1, -1)


            neg_x = tf.nn.embedding_lookup(self.item_emb, self.neg_p)#[B,neg,e]
            neg_x = tf.clip_by_norm(neg_x, 1, -1)

            self.pos_y = pos_y = self.pos_object_function(u, pos_v, m, pos_x, self.w) #[B,tsl]
            self.neg_y = neg_y = self.neg_object_function(u, neg_v, m, neg_x, self.w) #[B,nsl]

            #margin based hinge loss
            self.loss = self.loss_function(self.gamma, pos_y, neg_y)
            self.next_items = m
//  [python]
def log_training_summaries(self, tensor, name, family):
        """
        Make summaries for training. Plotting summaries for
        - Sparsity of tensor
        - Histogram of tensor (on log scale)
        - Denormals in tensor
        - Norm of tensor

        :param Tensor tensor: tensor to plot summaries for
        :param str name: name of the tensor to plot summaries for
        :param str family: family that the tensor belongs to (kernel / bias)
        """

        tf.compat.v1.summary.scalar(
            f"sparsity_{family}/{name}",
            (
                1.0
                - tf.math.count_nonzero(tensor, dtype=tf.float32)
                / tf.size(tensor, out_type=tf.float32)
            ),
        )
        tf.compat.v1.summary.scalar(
            f"denormal_{family}/{name}",
            (
                tf.reduce_sum(
                    tf.to_float(
                        tf.math.logical_and(
                            tf.math.less(tf.abs(tensor), self._denormal_range),
                            tf.math.not_equal(tensor, 0),
                        )
                    )
                )
                / tf.size(tensor, out_type=tf.float32)
            ),
        )
        tf.compat.v1.summary.scalar(
            f"norm_{family}/{name}", tf.linalg.global_norm([tensor])
        )
        if self._log_hists:
            tf.compat.v1.summary.histogram(
                f"{family}/{name}",
                tf.math.log(tf.to_float(tf.abs(tensor)) + 2.0 ** -50)
                / tf.math.log(2.0),
            )
//  [python]
def log_training_summaries(self, tensor, name, family):
        """
        Make summaries for training. Plotting summaries for
        - Sparsity of tensor
        - Histogram of tensor (on log scale)
        - Denormals in tensor
        - Norm of tensor

        :param Tensor tensor: tensor to plot summaries for
        :param str name: name of the tensor to plot summaries for
        :param str family: family that the tensor belongs to (kernel / bias)
        """

        tf.compat.v1.summary.scalar(
            f"sparsity_{family}/{name}",
            (
                1.0
                - tf.math.count_nonzero(tensor, dtype=tf.float32)
                / tf.size(tensor, out_type=tf.float32)
            ),
        )
        tf.compat.v1.summary.scalar(
            f"denormal_{family}/{name}",
            (
                tf.reduce_sum(
                    tf.cast(
                        tf.math.logical_and(
                            tf.math.less(tf.abs(tensor), self._denormal_range),
                            tf.math.not_equal(tensor, 0),
                        ),
                        tf.float32,
                    )
                )
                / tf.size(tensor, out_type=tf.float32)
            ),
        )
        tf.compat.v1.summary.scalar(
            f"norm_{family}/{name}", tf.linalg.global_norm([tensor])
        )
        if self._log_hists:
            tf.compat.v1.summary.histogram(
                f"{family}/{name}",
                tf.math.log(tf.cast(tf.abs(tensor), tf.float32) + 2.0 ** -50)
                / tf.math.log(2.0),
            )
//  [python]
def compute_sentence_correct(correct, sequence_mask):
    any_label = tf.reduce_max(tf.cast(sequence_mask, tf.int32), 1)
    sentence_correct_total = tf.reduce_sum(any_label)
    # is 1 when all is correct, 0 otherwise
    sentence_correct = tf.reduce_sum(tf.reduce_prod(
        tf.cast(
            tf.logical_or(correct, tf.logical_not(sequence_mask)),
            tf.int32
        ),
        1
    ) * any_label)
    return sentence_correct, sentence_correct_total
//  [python]
def build_train_ops(self, loss):
        """
        Setup optimizer and build train ops.

        :param Tensor loss: The loss tensor
        :return: Train ops
        """
        self._optimizer = self.build_optimizer()
        grads_and_vars = self._optimizer.compute_gradients(
            tf.cast(loss, tf.float32)
        )

        if self._log_summaries:
            self._gradient_global_norm = tf.linalg.global_norm(
                [g for (g, v) in grads_and_vars]
            )
            if not self.is_grad_accum():
                tf.compat.v1.summary.scalar(
                    'train/unclipped_grad_norm', self._gradient_global_norm
                )

        # This code mimics the CS1 dynamic loss scaling
        # kernel implementation, where the global norm of
        # weight gradients is computed to detect NaN/Inf and
        # is reused in gradient clipping. This saves compute
        # and simplifies kernel matching.
        if isinstance(
            self._optimizer,
            (LossScaleOptimizer, MixedPrecisionLossScaleOptimizer),
        ):
            # CSDynamicLossScale checks the global norm of the weight
            # gradients to determine whether a NaN/Inf has occurred when its
            # update() method is called. If its global_norm field is set, then
            # it will just check that value instead of recomputing the norm.
            if hasattr(self._optimizer.loss_scale, 'global_norm'):
                if self._gradient_global_norm is None:
                    self._gradient_global_norm = tf.linalg.global_norm(
                        [g for (g, v) in grads_and_vars]
                    )
                self._optimizer.loss_scale.global_norm = (
                    self._gradient_global_norm
                )

        clipped_grads_and_vars = self.clip_gradients(
            grads_and_vars, global_norm=self._gradient_global_norm
        )

        train_op = self._optimizer.apply_gradients(
            clipped_grads_and_vars,
            tf.compat.v1.train.get_or_create_global_step(),
        )

        if self._log_summaries and self._log_grads and not self.is_grad_accum():
            # Log the scaled gradients
            for (g, v) in grads_and_vars:
                if "kernel" in v.name:
                    self.log_training_summaries(
                        self._rescale(g), v.name, f"kernel_grads"
                    )
                elif "bias" in v.name:
                    self.log_training_summaries(
                        self._rescale(g), v.name, f"bias_grads"
                    )

        return train_op
//  [python]
def _apply_sparse(self, grad, var):
        beta1_power = t.cast(self._beta1_power, var.dtype.base_dtype)
        beta2_power = t.cast(self._beta2_power, var.dtype.base_dtype)
        lr_t = t.cast(self._lr_t, var.dtype.base_dtype)
        beta1_t = t.cast(self._beta1_t, var.dtype.base_dtype)
        beta2_t = t.cast(self._beta2_t, var.dtype.base_dtype)
        epsilon_t = t.cast(self._epsilon_t, var.dtype.base_dtype)
        lr = (lr_t * t.sqrt(1 - beta2_power) / (1 - beta1_power))

        # m := beta1 * m + (1 - beta1) * g_t
        # We use a slightly different version of the moving-average update formula
        # that does a better job of handling concurrent lockless updates:
        # m -= (1 - beta1) * (m - g_t)
        m = self.get_slot(var, "m")
        m_t_delta = tf.gather(m, grad.indices) - grad.values
        m_t = tf.scatter_sub(m, grad.indices,
                                    (1 - beta1_t) * m_t_delta,
                                    use_locking=self._use_locking)

        # v := beta2 * v + (1 - beta2) * (g_t * g_t)
        # We reformulate the update as:
        # v -= (1 - beta2) * (v - g_t * g_t)
        v = self.get_slot(var, "v")
        v_t_delta = tf.gather(v, grad.indices) - tf.square(grad.values)
        v_t = tf.scatter_sub(v, grad.indices,
                                    (1 - beta2_t) * v_t_delta,
                                    use_locking=self._use_locking)

        # variable -= learning_rate * m_t / (epsilon_t + sqrt(v_t))
        m_t_slice = tf.gather(m_t, grad.indices)
        v_t_slice = tf.gather(v_t, grad.indices)
        denominator_slice = tf.sqrt(v_t_slice) + epsilon_t
        var_update = tf.scatter_sub(var, grad.indices,
                                    lr * m_t_slice / denominator_slice,
                                    use_locking=self._use_locking)
        return tf.group(var_update, m_t, v_t)
//  [python]
def _apply_sparse(self, grad, var):
        beta1_power = tf.cast(self._beta1_power, var.dtype.base_dtype)
        beta2_power = tf.cast(self._beta2_power, var.dtype.base_dtype)
        lr_t = tf.cast(self._lr_t, var.dtype.base_dtype)
        beta1_t = tf.cast(self._beta1_t, var.dtype.base_dtype)
        beta2_t = tf.cast(self._beta2_t, var.dtype.base_dtype)
        epsilon_t = tf.cast(self._epsilon_t, var.dtype.base_dtype)
        lr = (lr_t * tf.sqrt(1 - beta2_power) / (1 - beta1_power))

        # m := beta1 * m + (1 - beta1) * g_t
        # We use a slightly different version of the moving-average update formula
        # that does a better job of handling concurrent lockless updates:
        # m -= (1 - beta1) * (m - g_t)
        m = self.get_slot(var, "m")
        m_t_delta = tf.gather(m, grad.indices) - grad.values
        m_t = tf.scatter_sub(m, grad.indices,
                                    (1 - beta1_t) * m_t_delta,
                                    use_locking=self._use_locking)

        # v := beta2 * v + (1 - beta2) * (g_t * g_t)
        # We reformulate the update as:
        # v -= (1 - beta2) * (v - g_t * g_t)
        v = self.get_slot(var, "v")
        v_t_delta = tf.gather(v, grad.indices) - tf.square(grad.values)
        v_t = tf.scatter_sub(v, grad.indices,
                                    (1 - beta2_t) * v_t_delta,
                                    use_locking=self._use_locking)

        # variable -= learning_rate * m_t / (epsilon_t + sqrt(v_t))
        m_t_slice = tf.gather(m_t, grad.indices)
        v_t_slice = tf.gather(v_t, grad.indices)
        denominator_slice = tf.sqrt(v_t_slice) + epsilon_t
        var_update = tf.scatter_sub(var, grad.indices,
                                    lr * m_t_slice / denominator_slice,
                                    use_locking=self._use_locking)
        return tf.group(var_update, m_t, v_t)
//  [python]
def build_F_init_layer(self):
        if self.F_init is not None:
            t_nz_idx, s_nz_idx = np.where(self.F_init>0)
            t_nz_idx, s_nz_idx = np.sort(t_nz_idx), np.sort(s_nz_idx)
            self.F_init = self.F_init[t_nz_idx,:]
            self.F_init = self.F_init[:,s_nz_idx]
            self.F_init_row=self.F_init/np.sum(self.F_init, axis=1, keepdims=True)
            self.F_init_col=np.transpose(self.F_init/np.sum(self.F_init, axis=0, keepdims=True))
            # print('self.F_init', self.F_init)
            # print('self.F_init.shape', self.F_init.shape)
            # take care of the initial dictionary
            # F_init_count = np.count_nonzero(self.F_init)
            # print('F_init_count', F_init_count)
            F_init_row = tf.convert_to_tensor(self.F_init_row, dtype=tf.float32)
            F_init_col = tf.convert_to_tensor(self.F_init_col, dtype=tf.float32)

            self.init_vecs_list = [None, None]
            self.init_vecs_list[0] = self.vecs_list[0][s_nz_idx,:]
            self.init_vecs_list[1] = self.vecs_list[1][t_nz_idx,:]

            if USE_EUC_DIST:
                self.align_loss_row_init = euclidean_distance(self.trans_s_init, self.input_t_init) / tf.to_float(self.F_init.shape[0])
                self.align_loss_col_init = euclidean_distance(self.trans_t_init, self.input_s_init) / tf.to_float(self.F_init.shape[1])
            else:
                self.align_loss_row_init = 1 - cosine_similarity(tf.matmul(F_init_row, self.trans_s_init), self.input_t_init) / tf.to_float(self.F_init.shape[0])
                self.align_loss_col_init = 1 - cosine_similarity(tf.matmul(F_init_col, self.trans_t_init), self.input_s_init) / tf.to_float(self.F_init.shape[1])
            self.init_align_loss = (self.align_loss_row_init + self.align_loss_col_init)/2
        else:
            self.init_align_loss = tf.constant(0, dtype=tf.float32)
//  [python]
def build_F_init_layer(self):
        if self.F_init is not None:
            t_nz_idx, s_nz_idx = np.where(self.F_init>0)
            t_nz_idx, s_nz_idx = np.sort(t_nz_idx), np.sort(s_nz_idx)
            self.F_init = self.F_init[t_nz_idx,:]
            self.F_init = self.F_init[:,s_nz_idx]
            self.F_init_row=self.F_init/np.sum(self.F_init, axis=1, keepdims=True)
            self.F_init_col=np.transpose(self.F_init/np.sum(self.F_init, axis=0, keepdims=True))
            # print('self.F_init', self.F_init)
            # print('self.F_init.shape', self.F_init.shape)
            # take care of the initial dictionary
            # F_init_count = np.count_nonzero(self.F_init)
            # print('F_init_count', F_init_count)
            F_init_row = tf.convert_to_tensor(self.F_init_row, dtype=tf.float32)
            F_init_col = tf.convert_to_tensor(self.F_init_col, dtype=tf.float32)

            self.init_vecs_list = [None, None]
            self.init_vecs_list[0] = self.vecs_list[0][s_nz_idx,:]
            self.init_vecs_list[1] = self.vecs_list[1][t_nz_idx,:]

            if USE_EUC_DIST:
                self.align_loss_row_init = euclidean_distance(self.trans_s_init, self.input_t_init) / tf.cast(self.F_init.shape[0], dtype=tf.float32)
                self.align_loss_col_init = euclidean_distance(self.trans_t_init, self.input_s_init) / tf.cast(self.F_init.shape[1], dtype=tf.float32)
            else:
                self.align_loss_row_init = 1 - cosine_similarity(tf.matmul(F_init_row, self.trans_s_init), self.input_t_init) / tf.cast(self.F_init.shape[0], dtype=tf.float32)
                self.align_loss_col_init = 1 - cosine_similarity(tf.matmul(F_init_col, self.trans_t_init), self.input_s_init) / tf.cast(self.F_init.shape[1], dtype=tf.float32)
            self.init_align_loss = (self.align_loss_row_init + self.align_loss_col_init)/2
        else:
            self.init_align_loss = tf.constant(0, dtype=tf.float32)
//  [python]
def build_true_false_positives(decoded, mask_batch_major, labels_batch_major,
                               classes, equals_label):
    masked_equals_label = tf.logical_and(equals_label, mask_batch_major)

    # now for each class compute tp, fp, fn
    # [nclasses x batch x time]
    masked_per_class = tf.logical_and(
        tf.equal(labels_batch_major[None, :, :], tf.range(classes)[:, None, None]),
        mask_batch_major)

    # correct, and on label
    correct = tf.reduce_sum(tf.to_int32(tf.logical_and(masked_per_class, equals_label[None, :, :])),
        axis=[1, 2])
    # predicted a particular class
    guessed = tf.reduce_sum(tf.to_int32(tf.logical_and(tf.equal(decoded[None, :, :], tf.range(classes)[:, None, None]), mask_batch_major)),
        axis=[1, 2])
    total = tf.reduce_sum(tf.to_int32(masked_per_class), axis=[1, 2])
    tp, fp, fn = correct, guessed - correct, total - correct

    tf.add_to_collection(TRUE_POSITIVES, tp)
    tf.add_to_collection(FALSE_POSITIVES, fp)
    tf.add_to_collection(FALSE_NEGATIVES, fn)
//  [python]
def build_true_false_positives(decoded, mask_batch_major, labels_batch_major,
                               classes, equals_label):
    masked_equals_label = tf.logical_and(equals_label, mask_batch_major)

    # now for each class compute tp, fp, fn
    # [nclasses x batch x time]
    masked_per_class = tf.logical_and(
        tf.equal(labels_batch_major[None, :, :], tf.range(classes)[:, None, None]),
        mask_batch_major)

    # correct, and on label
    correct = tf.reduce_sum(tf.cast(tf.logical_and(masked_per_class, equals_label[None, :, :]), tf.int32),
        axis=[1, 2])
    # predicted a particular class
    guessed = tf.reduce_sum(tf.cast(tf.logical_and(tf.equal(decoded[None, :, :], tf.range(classes)[:, None, None]), mask_batch_major), tf.int32),
        axis=[1, 2])
    total = tf.reduce_sum(tf.cast(masked_per_class, tf.int32), axis=[1, 2])
    tp, fp, fn = correct, guessed - correct, total - correct

    tf.add_to_collection(TRUE_POSITIVES, tp)
    tf.add_to_collection(FALSE_POSITIVES, fp)
    tf.add_to_collection(FALSE_NEGATIVES, fn)
//  [python]
def crf_metrics(unary_scores, labels, transition_params, sequence_lengths,
                mask):
    """
    Computes CRF output metrics.
    Receives:
        unary_scores : batch-major order
        labels : batch-major order
        transition_params : nclasses x nclasses matrix.
        sequence_lengths : length of each time-sequence
        mask : batch-major example mask

    Returns:
        token_correct,
        token_correct_total,
        sentence_correct,
        sentence_correct_total
    """
    classes = unary_scores.get_shape()[-1].value
    decoded, scores = viterbi_decode(unary_scores,
                                     transition_params,
                                     sequence_lengths)

    tf.add_to_collection(UNARY_SCORES, unary_scores)
    tf.add_to_collection(DECODED, decoded)
    tf.add_to_collection(DECODED_SCORES, scores)

    equals_label = tf.equal(labels, decoded)
    token_correct = tf.reduce_sum(
        tf.to_int32(
            tf.logical_and(equals_label, mask)
        )
    )
    token_correct_total = tf.reduce_sum(tf.to_int32(mask))
    tf.add_to_collection(TOKEN_CORRECT, token_correct)
    tf.add_to_collection(TOKEN_CORRECT_TOTAL, token_correct_total)
    sentence_correct, _ = compute_sentence_correct(equals_label, mask)
    sentence_correct_total = tf.reduce_sum(tf.to_int32(mask[:, 0]))

    tf.add_to_collection(SENTENCE_CORRECT, sentence_correct)
    tf.add_to_collection(SENTENCE_CORRECT_TOTAL, sentence_correct_total)

    build_true_false_positives(decoded, mask, labels,
        classes, equals_label)

    return (token_correct, token_correct_total,
            sentence_correct, sentence_correct_total)
//  [python]
def crf_metrics(unary_scores, labels, transition_params, sequence_lengths,
                mask):
    """
    Computes CRF output metrics.
    Receives:
        unary_scores : batch-major order
        labels : batch-major order
        transition_params : nclasses x nclasses matrix.
        sequence_lengths : length of each time-sequence
        mask : batch-major example mask

    Returns:
        token_correct,
        token_correct_total,
        sentence_correct,
        sentence_correct_total
    """
    classes = unary_scores.get_shape()[-1].value
    decoded, scores = viterbi_decode(unary_scores,
                                     transition_params,
                                     sequence_lengths)

    tf.add_to_collection(UNARY_SCORES, unary_scores)
    tf.add_to_collection(DECODED, decoded)
    tf.add_to_collection(DECODED_SCORES, scores)

    equals_label = tf.equal(labels, decoded)
    token_correct = tf.reduce_sum(
        tf.cast(
            tf.logical_and(equals_label, mask),
            tf.int32
        )
    )
    token_correct_total = tf.reduce_sum(tf.cast(mask, tf.int32))
    tf.add_to_collection(TOKEN_CORRECT, token_correct)
    tf.add_to_collection(TOKEN_CORRECT_TOTAL, token_correct_total)
    sentence_correct, _ = compute_sentence_correct(equals_label, mask)
    sentence_correct_total = tf.reduce_sum(tf.cast(mask[:, 0], tf.int32))

    tf.add_to_collection(SENTENCE_CORRECT, sentence_correct)
    tf.add_to_collection(SENTENCE_CORRECT_TOTAL, sentence_correct_total)

    build_true_false_positives(decoded, mask, labels,
        classes, equals_label)

    return (token_correct, token_correct_total,
            sentence_correct, sentence_correct_total)
//  [python]
def embedding_lookup(inputs,
                     size,
                     dim,
                     dtype,
                     reuse=None,
                     mask_negative=False,
                     trainable=True,
                     place_on_cpu_if_big=True):
    """
    Construct an Embedding layer that gathers
    elements from a matrix with `size` rows,
    and `dim` features using the indices stored in `x`.

    Arguments:
    ----------
        inputs : tf.Tensor, of integer type
        size : int, how many symbols in the lookup table
        dim : int, how many columns per symbol.
        dtype : data type for the lookup table (e.g. tf.float32)
        reuse : bool, (default None) whether the lookup table
            was already used before (thus this is weight sharing).
        mask_negative : bool, (default False) should -1s in the
            lookup input indicate padding (e.g. no lookup),
            and thus should those values be masked out post-lookup.
        trainable : bool (default True), whether the parameters of
            this lookup table can be backpropagated into (e.g.
            for Glove word vectors that are fixed pre-trained, this
            can be set to False).
        place_on_cpu_if_big : bool, if matrix is big, store it on cpu.
    Returns:
    --------
        tf.Tensor, result of tf.nn.embedding_lookup(LookupTable, inputs)
    """
    W = get_embedding_lookup(size, dim, dtype, reuse, trainable=trainable)
    if mask_negative:
        embedded = tf.nn.embedding_lookup(W, tf.maximum(inputs, 0))
        null_mask = tf.expand_dims(
            tf.to_float(
                tf.not_equal(inputs, -1),
            ),
            -1
        )
        return embedded * null_mask
    else:
        return tf.nn.embedding_lookup(W, inputs)
//  [python]
def embedding_lookup(inputs,
                     size,
                     dim,
                     dtype,
                     reuse=None,
                     mask_negative=False,
                     trainable=True,
                     place_on_cpu_if_big=True):
    """
    Construct an Embedding layer that gathers
    elements from a matrix with `size` rows,
    and `dim` features using the indices stored in `x`.

    Arguments:
    ----------
        inputs : tf.Tensor, of integer type
        size : int, how many symbols in the lookup table
        dim : int, how many columns per symbol.
        dtype : data type for the lookup table (e.g. tf.float32)
        reuse : bool, (default None) whether the lookup table
            was already used before (thus this is weight sharing).
        mask_negative : bool, (default False) should -1s in the
            lookup input indicate padding (e.g. no lookup),
            and thus should those values be masked out post-lookup.
        trainable : bool (default True), whether the parameters of
            this lookup table can be backpropagated into (e.g.
            for Glove word vectors that are fixed pre-trained, this
            can be set to False).
        place_on_cpu_if_big : bool, if matrix is big, store it on cpu.
    Returns:
    --------
        tf.Tensor, result of tf.nn.embedding_lookup(LookupTable, inputs)
    """
    W = get_embedding_lookup(size, dim, dtype, reuse, trainable=trainable)
    if mask_negative:
        embedded = tf.nn.embedding_lookup(W, tf.maximum(inputs, 0))
        null_mask = tf.expand_dims(
            tf.cast(
                tf.not_equal(inputs, -1),
                dtype
            ),
            -1
        )
        return embedded * null_mask
    else:
        return tf.nn.embedding_lookup(W, inputs)
//  [python]
def softmax_metrics(unary_scores, labels, mask):
    """
    Compute softmax output stats for correct/accuracy per-token/per-sentence.
    Receive
        unary_scores : time-major
        labels : time-major
        mask : time-major
    Returns:
        token_correct,
        token_correct_total,
        sentence_correct,
        sentence_correct_total
    """
    classes = unary_scores.get_shape()[-1].value
    unary_scores_batch_major = tf.transpose(unary_scores, [1, 0, 2])
    labels_batch_major = tf.transpose(labels, [1, 0])
    mask_batch_major = tf.transpose(mask, [1, 0])
    decoded = tf.cast(tf.argmax(unary_scores_batch_major, 2), labels.dtype)
    unary_probs_batch_major = tf.nn.softmax(unary_scores_batch_major)
    scores = tf.reduce_max(unary_probs_batch_major, 2)

    tf.add_to_collection(UNARY_SCORES, unary_probs_batch_major)
    tf.add_to_collection(DECODED, decoded)
    tf.add_to_collection(DECODED_SCORES, scores)

    equals_label = tf.equal(decoded, labels_batch_major)

    token_correct = tf.reduce_sum(
        tf.to_int32(
            tf.logical_and(
                equals_label,
                mask_batch_major
            )
        )
    )
    token_correct_total = tf.reduce_sum(tf.to_int32(mask))
    tf.add_to_collection(TOKEN_CORRECT, token_correct)
    tf.add_to_collection(TOKEN_CORRECT_TOTAL, token_correct_total)

    sentence_correct, sentence_correct_total = compute_sentence_correct(
        equals_label, mask_batch_major
    )
    tf.add_to_collection(SENTENCE_CORRECT, sentence_correct)
    tf.add_to_collection(SENTENCE_CORRECT_TOTAL, sentence_correct_total)

    build_true_false_positives(decoded, mask_batch_major, labels_batch_major,
        classes, equals_label)
    return (token_correct, token_correct_total,
            sentence_correct, sentence_correct_total)
//  [python]
def softmax_metrics(unary_scores, labels, mask):
    """
    Compute softmax output stats for correct/accuracy per-token/per-sentence.
    Receive
        unary_scores : time-major
        labels : time-major
        mask : time-major
    Returns:
        token_correct,
        token_correct_total,
        sentence_correct,
        sentence_correct_total
    """
    classes = unary_scores.get_shape()[-1].value
    unary_scores_batch_major = tf.transpose(unary_scores, [1, 0, 2])
    labels_batch_major = tf.transpose(labels, [1, 0])
    mask_batch_major = tf.transpose(mask, [1, 0])
    decoded = tf.cast(tf.argmax(unary_scores_batch_major, 2), labels.dtype)
    unary_probs_batch_major = tf.nn.softmax(unary_scores_batch_major)
    scores = tf.reduce_max(unary_probs_batch_major, 2)

    tf.add_to_collection(UNARY_SCORES, unary_probs_batch_major)
    tf.add_to_collection(DECODED, decoded)
    tf.add_to_collection(DECODED_SCORES, scores)

    equals_label = tf.equal(decoded, labels_batch_major)

    token_correct = tf.reduce_sum(
        tf.cast(
            tf.logical_and(
                equals_label,
                mask_batch_major
            ),
            tf.int32
        )
    )
    token_correct_total = tf.reduce_sum(tf.cast(mask, tf.int32))
    tf.add_to_collection(TOKEN_CORRECT, token_correct)
    tf.add_to_collection(TOKEN_CORRECT_TOTAL, token_correct_total)

    sentence_correct, sentence_correct_total = compute_sentence_correct(
        equals_label, mask_batch_major
    )
    tf.add_to_collection(SENTENCE_CORRECT, sentence_correct)
    tf.add_to_collection(SENTENCE_CORRECT_TOTAL, sentence_correct_total)

    build_true_false_positives(decoded, mask_batch_major, labels_batch_major,
        classes, equals_label)
    return (token_correct, token_correct_total,
            sentence_correct, sentence_correct_total)
//  [python]
def prepare_data_for_segmentation_loss(y_true, y_pred, num_classes=10, ignore_value=255):
    """
    Prepare predicted logits and ground-truth maps for the loss, removing pixels from ignored classes.
    :param y_true:        Ground-truth label map(s) (e.g., of shape B x H x W)
    :param y_pred:        Predicted logit map(s) () (e.g., of shape B x H x W x N, N number of classes)
    :param num_classes:   Number of classes
    :param ignore_value:  trainID value of ignored classes (`None` if ignored none) 
    :return:              Tensors edited, ready for the loss computation
    """

    with tf.name_scope('prepare_data_for_loss'):
        # Flattening the tensors to simplify the following operations:
        if len(y_pred.shape) > (len(y_true.shape) - 1):
            y_pred = tf.reshape(y_pred, [-1, num_classes])
        else:
            y_pred = tf.reshape(y_pred, [-1])
        y_true = tf.reshape(tf.to_int32(y_true), [-1])

        if ignore_value is not None:
            # To compare only on the considered class, we remove all the elements in the images
            # belonging to the ignored ones.
            # For that, we first compute the mask of the pixels belonging to valid labels:
            mask_for_valid_labels = get_mask_for_valid_labels(
                y_true, num_classes, ignore_value=ignore_value)
    
            # Then we use this mask to remove all pixels/elements not belonging to valid classes:
            y_true = tf.boolean_mask(y_true, mask_for_valid_labels, axis=0, name='gt_valid')
            y_pred = tf.boolean_mask(y_pred, mask_for_valid_labels, axis=0, name='pred_valid')

    return y_true, y_pred