jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _pad_to_window_size(
        self,
        input_ids,
        attention_mask,
        token_type_ids,
        position_ids,
        inputs_embeds,
        pad_token_id,
    ):
        """A helper function to pad tokens and mask to work with implementation of Longformer selfattention."""
        # padding
        attention_window = (
            self.attention_window if isinstance(self.attention_window, int) else max(self.attention_window)
        )

        assert attention_window % 2 == 0, f"`attention_window` should be an even value. Given {attention_window}"

        input_shape = shape_list(input_ids) if input_ids is not None else shape_list(inputs_embeds)
        batch_size, seq_len = input_shape[:2]
        padding_len = (attention_window - seq_len % attention_window) % attention_window

        paddings = tf.convert_to_tensor([[0, 0], [0, padding_len]])

        if input_ids is not None:
            input_ids = tf.pad(input_ids, paddings, constant_values=pad_token_id)

        if position_ids is not None:
            # pad with position_id = pad_token_id as in modeling_roberta.RobertaEmbeddings
            position_ids = tf.pad(position_ids, paddings, constant_values=pad_token_id)

        if inputs_embeds is not None:
            if padding_len > 0:
                input_ids_padding = tf.cast(tf.fill((batch_size, padding_len), self.pad_token_id), tf.int64)
                inputs_embeds_padding = self.embeddings(input_ids_padding)
                inputs_embeds = tf.concat([inputs_embeds, inputs_embeds_padding], axis=-2)

        attention_mask = tf.pad(attention_mask, paddings, constant_values=False)  # no attention on the padding tokens
        token_type_ids = tf.pad(token_type_ids, paddings, constant_values=0)  # pad with token_type_id = 0

        return (
            padding_len,
            input_ids,
            attention_mask,
            token_type_ids,
            position_ids,
            inputs_embeds,
        )
//  [python]
def path_rnn_last_state(self, is_evaluating, path_embed, path_lengths, valid_contexts_mask):
        # path_embed:           (batch, max_contexts, max_path_length+1, dim)
        # path_length:          (batch, max_contexts)
        # valid_contexts_mask:  (batch, max_contexts)
        max_contexts = tf.shape(path_embed)[1]
        flat_paths = tf.reshape(path_embed, shape=[-1, self.config.MAX_PATH_LENGTH,
                                                   self.config.EMBEDDINGS_SIZE])  # (batch * max_contexts, max_path_length+1, dim)
        flat_valid_contexts_mask = tf.reshape(valid_contexts_mask, [-1])  # (batch * max_contexts)
        lengths = tf.multiply(tf.reshape(path_lengths, [-1]),
                              tf.to_int32(flat_valid_contexts_mask))  # (batch * max_contexts)
        if self.config.BIRNN:
            rnn_cell_fw = tf.nn.rnn_cell.LSTMCell(self.config.RNN_SIZE / 2)
            rnn_cell_bw = tf.nn.rnn_cell.LSTMCell(self.config.RNN_SIZE / 2)
            if not is_evaluating:
                rnn_cell_fw = tf.nn.rnn_cell.DropoutWrapper(rnn_cell_fw,
                                                            output_keep_prob=self.config.RNN_DROPOUT_KEEP_PROB)
                rnn_cell_bw = tf.nn.rnn_cell.DropoutWrapper(rnn_cell_bw,
                                                            output_keep_prob=self.config.RNN_DROPOUT_KEEP_PROB)
            _, (state_fw, state_bw) = tf.nn.bidirectional_dynamic_rnn(
                cell_fw=rnn_cell_fw,
                cell_bw=rnn_cell_bw,
                inputs=flat_paths,
                dtype=tf.float32,
                sequence_length=lengths)
            final_rnn_state = tf.concat([state_fw.h, state_bw.h], axis=-1)  # (batch * max_contexts, rnn_size)  
        else:
            rnn_cell = tf.nn.rnn_cell.LSTMCell(self.config.RNN_SIZE)
            if not is_evaluating:
                rnn_cell = tf.nn.rnn_cell.DropoutWrapper(rnn_cell, output_keep_prob=self.config.RNN_DROPOUT_KEEP_PROB)
            _, state = tf.nn.dynamic_rnn(
                cell=rnn_cell,
                inputs=flat_paths,
                dtype=tf.float32,
                sequence_length=lengths
            )
            final_rnn_state = state.h  # (batch * max_contexts, rnn_size)

        return tf.reshape(final_rnn_state,
                          shape=[-1, max_contexts, self.config.RNN_SIZE])
//  [python]
def path_rnn_last_state(self, is_evaluating, path_embed, path_lengths, valid_contexts_mask):
        # path_embed:           (batch, max_contexts, max_path_length+1, dim)
        # path_length:          (batch, max_contexts)
        # valid_contexts_mask:  (batch, max_contexts)
        max_contexts = tf.shape(path_embed)[1]
        flat_paths = tf.reshape(path_embed, shape=[-1, self.config.MAX_PATH_LENGTH,
                                                   self.config.EMBEDDINGS_SIZE])  # (batch * max_contexts, max_path_length+1, dim)
        flat_valid_contexts_mask = tf.reshape(valid_contexts_mask, [-1])  # (batch * max_contexts)
        lengths = tf.multiply(tf.reshape(path_lengths, [-1]),
                              tf.cast(flat_valid_contexts_mask, tf.int32))  # (batch * max_contexts)
        if self.config.BIRNN:
            rnn_cell_fw = tf.nn.rnn_cell.LSTMCell(self.config.RNN_SIZE / 2)
            rnn_cell_bw = tf.nn.rnn_cell.LSTMCell(self.config.RNN_SIZE / 2)
            if not is_evaluating:
                rnn_cell_fw = tf.nn.rnn_cell.DropoutWrapper(rnn_cell_fw,
                                                            output_keep_prob=self.config.RNN_DROPOUT_KEEP_PROB)
                rnn_cell_bw = tf.nn.rnn_cell.DropoutWrapper(rnn_cell_bw,
                                                            output_keep_prob=self.config.RNN_DROPOUT_KEEP_PROB)
            _, (state_fw, state_bw) = tf.nn.bidirectional_dynamic_rnn(
                cell_fw=rnn_cell_fw,
                cell_bw=rnn_cell_bw,
                inputs=flat_paths,
                dtype=tf.float32,
                sequence_length=lengths)
            final_rnn_state = tf.concat([state_fw.h, state_bw.h], axis=-1)  # (batch * max_contexts, rnn_size)  
        else:
            rnn_cell = tf.nn.rnn_cell.LSTMCell(self.config.RNN_SIZE)
            if not is_evaluating:
                rnn_cell = tf.nn.rnn_cell.DropoutWrapper(rnn_cell, output_keep_prob=self.config.RNN_DROPOUT_KEEP_PROB)
            _, state = tf.nn.dynamic_rnn(
                cell=rnn_cell,
                inputs=flat_paths,
                dtype=tf.float32,
                sequence_length=lengths
            )
            final_rnn_state = state.h  # (batch * max_contexts, rnn_size)

        return tf.reshape(final_rnn_state,
                          shape=[-1, max_contexts, self.config.RNN_SIZE])
//  [python]
def call(
        self,
        input_ids=None,
        attention_mask=None,
        head_mask=None,
        inputs_embeds=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
        training=False,
    ):
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            input_shape = shape_list(input_ids)
        elif inputs_embeds is not None:
            input_shape = shape_list(inputs_embeds)[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        if attention_mask is None:
            attention_mask = tf.ones(input_shape)  # (bs, seq_length)

        attention_mask = tf.cast(attention_mask, dtype=tf.float32)

        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
        if head_mask is not None:
            raise NotImplementedError
        else:
            head_mask = [None] * self.num_hidden_layers

        embedding_output = self.embeddings(input_ids, inputs_embeds=inputs_embeds)  # (bs, seq_length, dim)
        tfmr_output = self.transformer(
            embedding_output,
            attention_mask,
            head_mask,
            output_attentions,
            output_hidden_states,
            return_dict,
            training=training,
        )

        return tfmr_output
//  [python]
def call(self, inputs_with_item_idx, training=None, **kwargs):
        user_vec, item_vec, item_idx = inputs_with_item_idx
        if item_idx.dtype != tf.int64:
            item_idx = tf.to_int64(item_idx)
        user_vec /= self.temperature
        logits = tf.matmul(user_vec, item_vec, transpose_b=True)
        loss = inbatch_softmax_cross_entropy_with_logits(logits, self.item_count, item_idx)
        return tf.expand_dims(loss, axis=1)
//  [python]
def call(self, inputs_with_item_idx, training=None, **kwargs):
        user_vec, item_vec, item_idx = inputs_with_item_idx
        if item_idx.dtype != tf.int64:
            item_idx = tf.cast(item_idx, tf.int64)
        user_vec /= self.temperature
        logits = tf.matmul(user_vec, item_vec, transpose_b=True)
        loss = inbatch_softmax_cross_entropy_with_logits(logits, self.item_count, item_idx)
        return tf.expand_dims(loss, axis=1)
//  [python]
def call(self, query, key, value, mask, head_mask, output_attentions, training=False):
        """
        Parameters:
            query: tf.Tensor(bs, seq_length, dim)
            key: tf.Tensor(bs, seq_length, dim)
            value: tf.Tensor(bs, seq_length, dim)
            mask: tf.Tensor(bs, seq_length)

        Returns:
            weights: tf.Tensor(bs, n_heads, seq_length, seq_length) Attention weights context: tf.Tensor(bs,
            seq_length, dim) Contextualized layer. Optional: only if `output_attentions=True`
        """
        bs, q_length, dim = shape_list(query)
        k_length = shape_list(key)[1]
        # assert dim == self.dim, f'Dimensions do not match: {dim} input vs {self.dim} configured'
        # assert key.size() == value.size()
        dim_per_head = int(self.dim / self.n_heads)
        dim_per_head = t.cast(dim_per_head, dtype=t.int32)
        mask_reshape = [bs, 1, 1, k_length]

        def shape(x):
            """separate heads"""
            return tf.transpose(tf.reshape(x, (bs, -1, self.n_heads, dim_per_head)), perm=(0, 2, 1, 3))

        def unshape(x):
            """group heads"""
            return tf.reshape(tf.transpose(x, perm=(0, 2, 1, 3)), (bs, -1, self.n_heads * dim_per_head))

        q = shape(self.q_lin(query))  # (bs, n_heads, q_length, dim_per_head)
        k = shape(self.k_lin(key))  # (bs, n_heads, k_length, dim_per_head)
        v = shape(self.v_lin(value))  # (bs, n_heads, k_length, dim_per_head)
        q = t.cast(q, dtype=t.float32)
        q = tf.multiply(q, tf.math.rsqrt(t.cast(dim_per_head, dtype=t.float32)))
        k = t.cast(k, dtype=q.dtype)
        scores = tf.matmul(q, k, transpose_b=True)  # (bs, n_heads, q_length, k_length)
        mask = tf.reshape(mask, mask_reshape)  # (bs, n_heads, qlen, klen)
        # scores.masked_fill_(mask, -float('inf'))            # (bs, n_heads, q_length, k_length)

        mask = t.cast(mask, dtype=scores.dtype)
        scores = scores - 1e30 * (1.0 - mask)
        weights = stable_softmax(scores, axis=-1)  # (bs, n_heads, qlen, klen)
        weights = self.dropout(weights, training=training)  # (bs, n_heads, qlen, klen)

        # Mask heads if we want to
        if head_mask is not None:
            weights = weights * head_mask

        context = tf.matmul(weights, v)  # (bs, n_heads, qlen, dim_per_head)
        context = unshape(context)  # (bs, q_length, dim)
        context = self.out_lin(context)  # (bs, q_length, dim)

        if output_attentions:
            return (context, weights)
        else:
            return (context,)
//  [python]
def call(self, query, key, value, mask, head_mask, output_attentions, training=False):
        """
        Parameters:
            query: tf.Tensor(bs, seq_length, dim)
            key: tf.Tensor(bs, seq_length, dim)
            value: tf.Tensor(bs, seq_length, dim)
            mask: tf.Tensor(bs, seq_length)

        Returns:
            weights: tf.Tensor(bs, n_heads, seq_length, seq_length) Attention weights context: tf.Tensor(bs,
            seq_length, dim) Contextualized layer. Optional: only if `output_attentions=True`
        """
        bs, q_length, dim = shape_list(query)
        k_length = shape_list(key)[1]
        # assert dim == self.dim, f'Dimensions do not match: {dim} input vs {self.dim} configured'
        # assert key.size() == value.size()
        dim_per_head = int(self.dim / self.n_heads)
        dim_per_head = tf.cast(dim_per_head, dtype=tf.int32)
        mask_reshape = [bs, 1, 1, k_length]

        def shape(x):
            """separate heads"""
            return tf.transpose(tf.reshape(x, (bs, -1, self.n_heads, dim_per_head)), perm=(0, 2, 1, 3))

        def unshape(x):
            """group heads"""
            return tf.reshape(tf.transpose(x, perm=(0, 2, 1, 3)), (bs, -1, self.n_heads * dim_per_head))

        q = shape(self.q_lin(query))  # (bs, n_heads, q_length, dim_per_head)
        k = shape(self.k_lin(key))  # (bs, n_heads, k_length, dim_per_head)
        v = shape(self.v_lin(value))  # (bs, n_heads, k_length, dim_per_head)
        q = tf.cast(q, dtype=tf.float32)
        q = tf.multiply(q, tf.math.rsqrt(tf.cast(dim_per_head, dtype=tf.float32)))
        k = tf.cast(k, dtype=q.dtype)
        scores = tf.matmul(q, k, transpose_b=True)  # (bs, n_heads, q_length, k_length)
        mask = tf.reshape(mask, mask_reshape)  # (bs, n_heads, qlen, klen)
        # scores.masked_fill_(mask, -float('inf'))            # (bs, n_heads, q_length, k_length)

        mask = tf.cast(mask, dtype=scores.dtype)
        scores = scores - 1e30 * (1.0 - mask)
        weights = stable_softmax(scores, axis=-1)  # (bs, n_heads, qlen, klen)
        weights = self.dropout(weights, training=training)  # (bs, n_heads, qlen, klen)

        # Mask heads if we want to
        if head_mask is not None:
            weights = weights * head_mask

        context = tf.matmul(weights, v)  # (bs, n_heads, qlen, dim_per_head)
        context = unshape(context)  # (bs, q_length, dim)
        context = self.out_lin(context)  # (bs, q_length, dim)

        if output_attentions:
            return (context, weights)
        else:
            return (context,)
//  [python]
def crf_filtered_inputs(inputs,
                        tag_bitmap):
  """Constrains the inputs to filter out certain tags at each time step.

  tag_bitmap limits the allowed tags at each input time step.
  This is useful when an observed output at a given time step needs to be
  constrained to a selected set of tags.

  Args:
    inputs: A [batch_size, max_seq_len, num_tags] tensor of unary potentials
      to use as input to the CRF layer.
    tag_bitmap: A [batch_size, max_seq_len, num_tags] boolean tensor
      representing all active tags at each index for which to calculate the
      unnormalized score.

  Returns:
    filtered_inputs: A [batch_size] vector of unnormalized sequence scores.
  """

  # set scores of filtered out inputs to be -inf.
  filtered_inputs = tf.where(
      tag_bitmap,
      inputs,
      tf.fill(tf.shape(inputs), t.float32(-inf)),
  )
  return filtered_inputs
//  [python]
def crf_filtered_inputs(inputs,
                        tag_bitmap):
  """Constrains the inputs to filter out certain tags at each time step.

  tag_bitmap limits the allowed tags at each input time step.
  This is useful when an observed output at a given time step needs to be
  constrained to a selected set of tags.

  Args:
    inputs: A [batch_size, max_seq_len, num_tags] tensor of unary potentials
      to use as input to the CRF layer.
    tag_bitmap: A [batch_size, max_seq_len, num_tags] boolean tensor
      representing all active tags at each index for which to calculate the
      unnormalized score.

  Returns:
    filtered_inputs: A [batch_size] vector of unnormalized sequence scores.
  """

  # set scores of filtered out inputs to be -inf.
  filtered_inputs = tf.where(
      tag_bitmap,
      inputs,
      tf.fill(tf.shape(inputs), tf.cast(float("-inf"), inputs.dtype)),
  )
  return filtered_inputs
//  [python]
def call(self, inputs_with_item_idx, training=None, **kwargs):
        item_embeddings, user_vec, item_idx = inputs_with_item_idx
        if item_idx.dtype != tf.int64:
            item_idx = tf.to_int64(item_idx)
        user_vec /= self.temperature
        if self.sampler == "inbatch":
            item_vec = tf.gather(item_embeddings, tf.squeeze(item_idx, axis=1))
            logits = tf.matmul(user_vec, item_vec, transpose_b=True)
            loss = inbatch_softmax_cross_entropy_with_logits(logits, self.item_count, item_idx)

        else:
            num_sampled = self.sampler_config['num_sampled']
            if self.sampler == "frequency":
                sampled_values = tf.nn.fixed_unigram_candidate_sampler(item_idx, 1, num_sampled, True,
                                                                       self.vocabulary_size,
                                                                       distortion=self.sampler_config['distortion'],
                                                                       unigrams=np.maximum(self.item_count, 1).tolist(),
                                                                       seed=None,
                                                                       name=None)
            elif self.sampler == "adaptive":
                sampled_values = tf.nn.learned_unigram_candidate_sampler(item_idx, 1, num_sampled, True,
                                                                         self.vocabulary_size, seed=None, name=None)
            elif self.sampler == "uniform":
                try:
                    sampled_values = tf.nn.uniform_candidate_sampler(item_idx, 1, num_sampled, True,
                                                                     self.vocabulary_size, seed=None, name=None)
                except AttributeError:
                    sampled_values = tf.random.uniform_candidate_sampler(item_idx, 1, num_sampled, True,
                                                                         self.vocabulary_size, seed=None, name=None)
            else:
                raise ValueError(' `%s` sampler is not supported ' % self.sampler)

            loss = tf.nn.sampled_softmax_loss(weights=item_embeddings,
                                              biases=self.zero_bias,
                                              labels=item_idx,
                                              inputs=user_vec,
                                              num_sampled=num_sampled,
                                              num_classes=self.vocabulary_size,
                                              sampled_values=sampled_values
                                              )
        return tf.expand_dims(loss, axis=1)
//  [python]
def call(self, inputs_with_item_idx, training=None, **kwargs):
        item_embeddings, user_vec, item_idx = inputs_with_item_idx
        if item_idx.dtype != tf.int64:
            item_idx = tf.cast(item_idx, tf.int64)
        user_vec /= self.temperature
        if self.sampler == "inbatch":
            item_vec = tf.gather(item_embeddings, tf.squeeze(item_idx, axis=1))
            logits = tf.matmul(user_vec, item_vec, transpose_b=True)
            loss = inbatch_softmax_cross_entropy_with_logits(logits, self.item_count, item_idx)

        else:
            num_sampled = self.sampler_config['num_sampled']
            if self.sampler == "frequency":
                sampled_values = tf.nn.fixed_unigram_candidate_sampler(item_idx, 1, num_sampled, True,
                                                                       self.vocabulary_size,
                                                                       distortion=self.sampler_config['distortion'],
                                                                       unigrams=np.maximum(self.item_count, 1).tolist(),
                                                                       seed=None,
                                                                       name=None)
            elif self.sampler == "adaptive":
                sampled_values = tf.nn.learned_unigram_candidate_sampler(item_idx, 1, num_sampled, True,
                                                                         self.vocabulary_size, seed=None, name=None)
            elif self.sampler == "uniform":
                try:
                    sampled_values = tf.nn.uniform_candidate_sampler(item_idx, 1, num_sampled, True,
                                                                     self.vocabulary_size, seed=None, name=None)
                except AttributeError:
                    sampled_values = tf.random.uniform_candidate_sampler(item_idx, 1, num_sampled, True,
                                                                         self.vocabulary_size, seed=None, name=None)
            else:
                raise ValueError(' `%s` sampler is not supported ' % self.sampler)

            loss = tf.nn.sampled_softmax_loss(weights=item_embeddings,
                                              biases=self.zero_bias,
                                              labels=item_idx,
                                              inputs=user_vec,
                                              num_sampled=num_sampled,
                                              num_classes=self.vocabulary_size,
                                              sampled_values=sampled_values
                                              )
        return tf.expand_dims(loss, axis=1)
//  [python]
def crf_sequence_score(
    inputs,
    tag_indices,
    sequence_lengths,
    transition_params,
):
  """Computes the unnormalized score for a tag sequence.

  Args:
    inputs: A [batch_size, max_seq_len, num_tags] tensor of unary potentials
      to use as input to the CRF layer.
    tag_indices: A [batch_size, max_seq_len] matrix of tag indices for which
      we compute the unnormalized score.
    sequence_lengths: A [batch_size] vector of true sequence lengths.
    transition_params: A [num_tags, num_tags] transition matrix.

  Returns:
    sequence_scores: A [batch_size] vector of unnormalized sequence scores.
  """
  tag_indices = t.cast(tag_indices, dtype=t.int32)
  sequence_lengths = t.cast(sequence_lengths, dtype=t.int32)

  def _multi_seq_fn():
    # Compute the scores of the given tag sequence.
    unary_scores = crf_unary_score(tag_indices, sequence_lengths, inputs)
    binary_scores = crf_binary_score(tag_indices, sequence_lengths,
                                     transition_params)
    sequence_scores = unary_scores + binary_scores
    return sequence_scores

  return _multi_seq_fn()
//  [python]
def crf_sequence_score(
    inputs,
    tag_indices,
    sequence_lengths,
    transition_params,
):
  """Computes the unnormalized score for a tag sequence.

  Args:
    inputs: A [batch_size, max_seq_len, num_tags] tensor of unary potentials
      to use as input to the CRF layer.
    tag_indices: A [batch_size, max_seq_len] matrix of tag indices for which
      we compute the unnormalized score.
    sequence_lengths: A [batch_size] vector of true sequence lengths.
    transition_params: A [num_tags, num_tags] transition matrix.

  Returns:
    sequence_scores: A [batch_size] vector of unnormalized sequence scores.
  """
  tag_indices = tf.cast(tag_indices, dtype=tf.int32)
  sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)

  def _multi_seq_fn():
    # Compute the scores of the given tag sequence.
    unary_scores = crf_unary_score(tag_indices, sequence_lengths, inputs)
    binary_scores = crf_binary_score(tag_indices, sequence_lengths,
                                     transition_params)
    sequence_scores = unary_scores + binary_scores
    return sequence_scores

  return _multi_seq_fn()
//  [python]
def crf_log_norm(inputs, sequence_lengths,
                 transition_params):
  """Computes the normalization for a CRF.

  Args:
    inputs: A [batch_size, max_seq_len, num_tags] tensor of unary potentials
      to use as input to the CRF layer.
    sequence_lengths: A [batch_size] vector of true sequence lengths.
    transition_params: A [num_tags, num_tags] transition matrix.

  Returns:
    log_norm: A [batch_size] vector of normalizers for a CRF.
  """
  sequence_lengths = t.cast(sequence_lengths, dtype=t.int32)
  # Split up the first and rest of the inputs in preparation for the forward
  # algorithm.
  first_input = tf.slice(inputs, [0, 0, 0], [-1, 1, -1])
  first_input = tf.squeeze(first_input, [1])

  def _multi_seq_fn():
    """Forward computation of alpha values."""
    rest_of_input = tf.slice(inputs, [0, 1, 0], [-1, -1, -1])
    # Compute the alpha values in the forward algorithm in order to get the
    # partition function.

    alphas = crf_forward(rest_of_input, first_input, transition_params,
                         sequence_lengths)
    log_norm = tf.reduce_logsumexp(alphas, [1])
    # Mask `log_norm` of the sequences with length <= zero.
    log_norm = tf.where(
        tf.less_equal(sequence_lengths, 0), tf.zeros_like(log_norm), log_norm)
    return log_norm

  return _multi_seq_fn()
//  [python]
def crf_log_norm(inputs, sequence_lengths,
                 transition_params):
  """Computes the normalization for a CRF.

  Args:
    inputs: A [batch_size, max_seq_len, num_tags] tensor of unary potentials
      to use as input to the CRF layer.
    sequence_lengths: A [batch_size] vector of true sequence lengths.
    transition_params: A [num_tags, num_tags] transition matrix.

  Returns:
    log_norm: A [batch_size] vector of normalizers for a CRF.
  """
  sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)
  # Split up the first and rest of the inputs in preparation for the forward
  # algorithm.
  first_input = tf.slice(inputs, [0, 0, 0], [-1, 1, -1])
  first_input = tf.squeeze(first_input, [1])

  def _multi_seq_fn():
    """Forward computation of alpha values."""
    rest_of_input = tf.slice(inputs, [0, 1, 0], [-1, -1, -1])
    # Compute the alpha values in the forward algorithm in order to get the
    # partition function.

    alphas = crf_forward(rest_of_input, first_input, transition_params,
                         sequence_lengths)
    log_norm = tf.reduce_logsumexp(alphas, [1])
    # Mask `log_norm` of the sequences with length <= zero.
    log_norm = tf.where(
        tf.less_equal(sequence_lengths, 0), tf.zeros_like(log_norm), log_norm)
    return log_norm

  return _multi_seq_fn()
//  [python]
def crf_multitag_sequence_score(
    inputs,
    tag_bitmap,
    sequence_lengths,
    transition_params,
):
  """Computes the unnormalized score of all tag sequences matching tag_bitmap.

  tag_bitmap enables more than one tag to be considered correct at each time
  step. This is useful when an observed output at a given time step is
  consistent with more than one tag, and thus the log likelihood of that
  observation must take into account all possible consistent tags.

  Using one-hot vectors in tag_bitmap gives results identical to
  crf_sequence_score.

  Args:
    inputs: A [batch_size, max_seq_len, num_tags] tensor of unary potentials
      to use as input to the CRF layer.
    tag_bitmap: A [batch_size, max_seq_len, num_tags] boolean tensor
      representing all active tags at each index for which to calculate the
      unnormalized score.
    sequence_lengths: A [batch_size] vector of true sequence lengths.
    transition_params: A [num_tags, num_tags] transition matrix.

  Returns:
    sequence_scores: A [batch_size] vector of unnormalized sequence scores.
  """
  tag_bitmap = t.cast(tag_bitmap, dtype=t.bool)
  sequence_lengths = t.cast(sequence_lengths, dtype=t.int32)
  filtered_inputs = crf_filtered_inputs(inputs, tag_bitmap)

  # If max_seq_len is 1, we skip the score calculation and simply gather the
  # unary potentials of all active tags.
  def _single_seq_fn():
    return tf.reduce_logsumexp(filtered_inputs, axis=[1, 2], keepdims=False)

  def _multi_seq_fn():
    # Compute the logsumexp of all scores of sequences
    # matching the given tags.
    return crf_log_norm(
        inputs=filtered_inputs,
        sequence_lengths=sequence_lengths,
        transition_params=transition_params,
    )

  return tf.cond(
      tf.equal(tf.shape(inputs)[1], 1), _single_seq_fn, _multi_seq_fn)
//  [python]
def crf_multitag_sequence_score(
    inputs,
    tag_bitmap,
    sequence_lengths,
    transition_params,
):
  """Computes the unnormalized score of all tag sequences matching tag_bitmap.

  tag_bitmap enables more than one tag to be considered correct at each time
  step. This is useful when an observed output at a given time step is
  consistent with more than one tag, and thus the log likelihood of that
  observation must take into account all possible consistent tags.

  Using one-hot vectors in tag_bitmap gives results identical to
  crf_sequence_score.

  Args:
    inputs: A [batch_size, max_seq_len, num_tags] tensor of unary potentials
      to use as input to the CRF layer.
    tag_bitmap: A [batch_size, max_seq_len, num_tags] boolean tensor
      representing all active tags at each index for which to calculate the
      unnormalized score.
    sequence_lengths: A [batch_size] vector of true sequence lengths.
    transition_params: A [num_tags, num_tags] transition matrix.

  Returns:
    sequence_scores: A [batch_size] vector of unnormalized sequence scores.
  """
  tag_bitmap = tf.cast(tag_bitmap, dtype=tf.bool)
  sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)
  filtered_inputs = crf_filtered_inputs(inputs, tag_bitmap)

  # If max_seq_len is 1, we skip the score calculation and simply gather the
  # unary potentials of all active tags.
  def _single_seq_fn():
    return tf.reduce_logsumexp(filtered_inputs, axis=[1, 2], keepdims=False)

  def _multi_seq_fn():
    # Compute the logsumexp of all scores of sequences
    # matching the given tags.
    return crf_log_norm(
        inputs=filtered_inputs,
        sequence_lengths=sequence_lengths,
        transition_params=transition_params,
    )

  return tf.cond(
      tf.equal(tf.shape(inputs)[1], 1), _single_seq_fn, _multi_seq_fn)
//  [python]
def call(self, inputs, **kwargs):

        behavior_embedding = inputs[0]
        seq_len = inputs[1]
        batch_size = tf.shape(behavior_embedding)[0]

        mask = tf.reshape(tf.sequence_mask(seq_len, self.max_len, tf.float32), [-1, self.max_len, 1, 1])

        behavior_embedding_mapping = tf.tensordot(behavior_embedding, self.bilinear_mapping_matrix, axes=1)
        behavior_embedding_mapping = tf.expand_dims(behavior_embedding_mapping, axis=2)

        behavior_embdding_mapping_ = tf.stop_gradient(behavior_embedding_mapping)  # N,max_len,1,E
        try:
            routing_logits = tf.truncated_normal([batch_size, self.max_len, self.k_max, 1], stddev=self.init_std)
        except AttributeError:
            routing_logits = tf.compat.v1.truncated_normal([batch_size, self.max_len, self.k_max, 1],
                                                           stddev=self.init_std)
        routing_logits = tf.stop_gradient(routing_logits)

        k_user = None
        if len(inputs) == 3:
            k_user = inputs[2]
            interest_mask = tf.sequence_mask(k_user, self.k_max, tf.float32)
            interest_mask = tf.reshape(interest_mask, [batch_size, 1, self.k_max, 1])
            interest_mask = tf.tile(interest_mask, [1, self.max_len, 1, 1])

            interest_padding = tf.ones_like(interest_mask) * -2 ** 31
            interest_mask = tf.to_float(interest_mask)

        for i in range(self.iteration_times):
            if k_user is not None:
                routing_logits = tf.where(interest_mask, routing_logits, interest_padding)
            try:
                weight = softmax(routing_logits, 2) * mask
            except TypeError:
                weight = tf.transpose(softmax(tf.transpose(routing_logits, [0, 1, 3, 2])),
                                      [0, 1, 3, 2]) * mask  # N,max_len,k_max,1
            if i < self.iteration_times - 1:
                Z = reduce_sum(tf.matmul(weight, behavior_embdding_mapping_), axis=1, keep_dims=True)  # N,1,k_max,E
                interest_capsules = squash(Z)
                delta_routing_logits = reduce_sum(
                    interest_capsules * behavior_embdding_mapping_,
                    axis=-1, keep_dims=True
                )
                routing_logits += delta_routing_logits
            else:
                Z = reduce_sum(tf.matmul(weight, behavior_embedding_mapping), axis=1, keep_dims=True)
                interest_capsules = squash(Z)

        interest_capsules = tf.reshape(interest_capsules, [-1, self.k_max, self.out_units])
        return interest_capsules
//  [python]
def call(self, inputs, **kwargs):

        behavior_embedding = inputs[0]
        seq_len = inputs[1]
        batch_size = tf.shape(behavior_embedding)[0]

        mask = tf.reshape(tf.sequence_mask(seq_len, self.max_len, tf.float32), [-1, self.max_len, 1, 1])

        behavior_embedding_mapping = tf.tensordot(behavior_embedding, self.bilinear_mapping_matrix, axes=1)
        behavior_embedding_mapping = tf.expand_dims(behavior_embedding_mapping, axis=2)

        behavior_embdding_mapping_ = tf.stop_gradient(behavior_embedding_mapping)  # N,max_len,1,E
        try:
            routing_logits = tf.truncated_normal([batch_size, self.max_len, self.k_max, 1], stddev=self.init_std)
        except AttributeError:
            routing_logits = tf.compat.v1.truncated_normal([batch_size, self.max_len, self.k_max, 1],
                                                           stddev=self.init_std)
        routing_logits = tf.stop_gradient(routing_logits)

        k_user = None
        if len(inputs) == 3:
            k_user = inputs[2]
            interest_mask = tf.sequence_mask(k_user, self.k_max, tf.float32)
            interest_mask = tf.reshape(interest_mask, [batch_size, 1, self.k_max, 1])
            interest_mask = tf.tile(interest_mask, [1, self.max_len, 1, 1])

            interest_padding = tf.ones_like(interest_mask) * -2 ** 31
            interest_mask = tf.cast(interest_mask, tf.bool)

        for i in range(self.iteration_times):
            if k_user is not None:
                routing_logits = tf.where(interest_mask, routing_logits, interest_padding)
            try:
                weight = softmax(routing_logits, 2) * mask
            except TypeError:
                weight = tf.transpose(softmax(tf.transpose(routing_logits, [0, 1, 3, 2])),
                                      [0, 1, 3, 2]) * mask  # N,max_len,k_max,1
            if i < self.iteration_times - 1:
                Z = reduce_sum(tf.matmul(weight, behavior_embdding_mapping_), axis=1, keep_dims=True)  # N,1,k_max,E
                interest_capsules = squash(Z)
                delta_routing_logits = reduce_sum(
                    interest_capsules * behavior_embdding_mapping_,
                    axis=-1, keep_dims=True
                )
                routing_logits += delta_routing_logits
            else:
                Z = reduce_sum(tf.matmul(weight, behavior_embedding_mapping), axis=1, keep_dims=True)
                interest_capsules = squash(Z)

        interest_capsules = tf.reshape(interest_capsules, [-1, self.k_max, self.out_units])
        return interest_capsules
//  [python]
def crf_log_likelihood(
    inputs,
    tag_indices,
    sequence_lengths,
    transition_params = None,
):
  """Computes the log-likelihood of tag sequences in a CRF.

  Args:
    inputs: A [batch_size, max_seq_len, num_tags] tensor of unary potentials
      to use as input to the CRF layer.
    tag_indices: A [batch_size, max_seq_len] matrix of tag indices for which
      we compute the log-likelihood.
    sequence_lengths: A [batch_size] vector of true sequence lengths.
    transition_params: A [num_tags, num_tags] transition matrix, if available.

  Returns:
    log_likelihood: A [batch_size] `Tensor` containing the log-likelihood of
      each example, given the sequence of tag indices.
    transition_params: A [num_tags, num_tags] transition matrix. This is
        either provided by the caller or created in this function.
  """
  inputs = tf.convert_to_tensor(inputs)

  num_tags = inputs.shape[2]

  # cast type to handle different types
  tag_indices = tf.to_int32(tag_indices)
  sequence_lengths = tf.to_int32(sequence_lengths)

  if transition_params is None:
    initializer = tf.keras.initializers.GlorotUniform()
    transition_params = tf.Variable(
        initializer([num_tags, num_tags]), "transitions")
  transition_params = tf.cast(transition_params, inputs.dtype)
  sequence_scores = crf_sequence_score(inputs, tag_indices, sequence_lengths,
                                       transition_params)
  log_norm = crf_log_norm(inputs, sequence_lengths, transition_params)

  # Normalize the scores to get the log-likelihood per example.
  log_likelihood = sequence_scores - log_norm
  return log_likelihood, transition_params
//  [python]
def crf_log_likelihood(
    inputs,
    tag_indices,
    sequence_lengths,
    transition_params = None,
):
  """Computes the log-likelihood of tag sequences in a CRF.

  Args:
    inputs: A [batch_size, max_seq_len, num_tags] tensor of unary potentials
      to use as input to the CRF layer.
    tag_indices: A [batch_size, max_seq_len] matrix of tag indices for which
      we compute the log-likelihood.
    sequence_lengths: A [batch_size] vector of true sequence lengths.
    transition_params: A [num_tags, num_tags] transition matrix, if available.

  Returns:
    log_likelihood: A [batch_size] `Tensor` containing the log-likelihood of
      each example, given the sequence of tag indices.
    transition_params: A [num_tags, num_tags] transition matrix. This is
        either provided by the caller or created in this function.
  """
  inputs = tf.convert_to_tensor(inputs)

  num_tags = inputs.shape[2]

  # cast type to handle different types
  tag_indices = tf.cast(tag_indices, dtype=tf.int32)
  sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)

  if transition_params is None:
    initializer = tf.keras.initializers.GlorotUniform()
    transition_params = tf.Variable(
        initializer([num_tags, num_tags]), "transitions")
  transition_params = tf.cast(transition_params, inputs.dtype)
  sequence_scores = crf_sequence_score(inputs, tag_indices, sequence_lengths,
                                       transition_params)
  log_norm = crf_log_norm(inputs, sequence_lengths, transition_params)

  # Normalize the scores to get the log-likelihood per example.
  log_likelihood = sequence_scores - log_norm
  return log_likelihood, transition_params
//  [python]
def crf_unary_score(tag_indices, sequence_lengths,
                    inputs):
  """Computes the unary scores of tag sequences.

  Args:
    tag_indices: A [batch_size, max_seq_len] matrix of tag indices.
    sequence_lengths: A [batch_size] vector of true sequence lengths.
    inputs: A [batch_size, max_seq_len, num_tags] tensor of unary potentials.

  Returns:
    unary_scores: A [batch_size] vector of unary scores.
  """
  tag_indices = t.cast(tag_indices, dtype=t.int32)
  sequence_lengths = t.cast(sequence_lengths, dtype=t.int32)

  batch_size = t.shape(inputs)[0]
  max_seq_len = t.shape(inputs)[1]
  num_tags = t.shape(inputs)[2]

  flattened_inputs = t.reshape(inputs, [-1])

  offsets = t.expand_dims(t.range(batch_size) * max_seq_len * num_tags, 1)
  offsets += t.expand_dims(t.range(max_seq_len) * num_tags, 0)
  # Use int32 or int64 based on tag_indices' dtype.
  if tag_indices.dtype == t.int64:
    offsets = t.cast(offsets, t.int64)
  flattened_tag_indices = t.reshape(offsets + tag_indices, [-1])

  unary_scores = t.reshape(
      t.gather(flattened_inputs, flattened_tag_indices),
      [batch_size, max_seq_len])

  masks = t.sequence_mask(
      sequence_lengths,
      maxlen=t.shape(tag_indices)[1],
      dtype=unary_scores.dtype)

  unary_scores = t.reduce_sum(unary_scores * masks, 1)
  return unary_scores
//  [python]
def crf_unary_score(tag_indices, sequence_lengths,
                    inputs):
  """Computes the unary scores of tag sequences.

  Args:
    tag_indices: A [batch_size, max_seq_len] matrix of tag indices.
    sequence_lengths: A [batch_size] vector of true sequence lengths.
    inputs: A [batch_size, max_seq_len, num_tags] tensor of unary potentials.

  Returns:
    unary_scores: A [batch_size] vector of unary scores.
  """
  tag_indices = tf.cast(tag_indices, dtype=tf.int32)
  sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)

  batch_size = tf.shape(inputs)[0]
  max_seq_len = tf.shape(inputs)[1]
  num_tags = tf.shape(inputs)[2]

  flattened_inputs = tf.reshape(inputs, [-1])

  offsets = tf.expand_dims(tf.range(batch_size) * max_seq_len * num_tags, 1)
  offsets += tf.expand_dims(tf.range(max_seq_len) * num_tags, 0)
  # Use int32 or int64 based on tag_indices' dtype.
  if tag_indices.dtype == tf.int64:
    offsets = tf.cast(offsets, tf.int64)
  flattened_tag_indices = tf.reshape(offsets + tag_indices, [-1])

  unary_scores = tf.reshape(
      tf.gather(flattened_inputs, flattened_tag_indices),
      [batch_size, max_seq_len])

  masks = tf.sequence_mask(
      sequence_lengths,
      maxlen=tf.shape(tag_indices)[1],
      dtype=unary_scores.dtype)

  unary_scores = tf.reduce_sum(unary_scores * masks, 1)
  return unary_scores
//  [python]
def call(self, inputs, state):
    """Build the CrfDecodeForwardRnnCell.

    Args:
      inputs: A [batch_size, num_tags] matrix of unary potentials.
      state: A [batch_size, num_tags] matrix containing the previous step's
        score values.

    Returns:
      backpointers: A [batch_size, num_tags] matrix of backpointers.
      new_state: A [batch_size, num_tags] matrix of new score values.
    """
    state = tf.expand_dims(state[0], 2)
    transition_scores = state + t.cast(self._transition_params,
                                        self._compute_dtype)
    new_state = inputs + tf.reduce_max(transition_scores, [1])
    backpointers = tf.argmax(transition_scores, 1)
    backpointers = tf.cast(backpointers, dtype=tf.int32)
    return backpointers, new_state
//  [python]
def call(self, inputs, state):
    """Build the CrfDecodeForwardRnnCell.

    Args:
      inputs: A [batch_size, num_tags] matrix of unary potentials.
      state: A [batch_size, num_tags] matrix containing the previous step's
        score values.

    Returns:
      backpointers: A [batch_size, num_tags] matrix of backpointers.
      new_state: A [batch_size, num_tags] matrix of new score values.
    """
    state = tf.expand_dims(state[0], 2)
    transition_scores = state + tf.cast(self._transition_params,
                                        self._compute_dtype)
    new_state = inputs + tf.reduce_max(transition_scores, [1])
    backpointers = tf.argmax(transition_scores, 1)
    backpointers = tf.cast(backpointers, dtype=tf.int32)
    return backpointers, new_state
//  [python]
def crf_binary_score(tag_indices, sequence_lengths,
                     transition_params):
  """Computes the binary scores of tag sequences.

  Args:
    tag_indices: A [batch_size, max_seq_len] matrix of tag indices.
    sequence_lengths: A [batch_size] vector of true sequence lengths.
    transition_params: A [num_tags, num_tags] matrix of binary potentials.

  Returns:
    binary_scores: A [batch_size] vector of binary scores.
  """
  tag_indices = t.cast(tag_indices, dtype=t.int32)
  sequence_lengths = t.cast(sequence_lengths, dtype=t.int32)

  num_tags = t.shape(transition_params)[0]
  num_transitions = t.shape(tag_indices)[1] - 1

  # Truncate by one on each side of the sequence to get the start and end
  # indices of each transition.
  start_tag_indices = t.slice(tag_indices, [0, 0], [-1, num_transitions])
  end_tag_indices = t.slice(tag_indices, [0, 1], [-1, num_transitions])

  # Encode the indices in a flattened representation.
  flattened_transition_indices = start_tag_indices * num_tags + end_tag_indices
  flattened_transition_params = t.reshape(transition_params, [-1])

  # Get the binary scores based on the flattened representation.
  binary_scores = t.gather(flattened_transition_params,
                            flattened_transition_indices)

  masks = t.sequence_mask(
      sequence_lengths,
      maxlen=t.shape(tag_indices)[1],
      dtype=binary_scores.dtype)
  truncated_masks = t.slice(masks, [0, 1], [-1, -1])
  binary_scores = t.reduce_sum(binary_scores * truncated_masks, 1)
  return binary_scores
//  [python]
def crf_binary_score(tag_indices, sequence_lengths,
                     transition_params):
  """Computes the binary scores of tag sequences.

  Args:
    tag_indices: A [batch_size, max_seq_len] matrix of tag indices.
    sequence_lengths: A [batch_size] vector of true sequence lengths.
    transition_params: A [num_tags, num_tags] matrix of binary potentials.

  Returns:
    binary_scores: A [batch_size] vector of binary scores.
  """
  tag_indices = tf.cast(tag_indices, dtype=tf.int32)
  sequence_lengths = tf.cast(sequence_lengths, dtype=tf.int32)

  num_tags = tf.shape(transition_params)[0]
  num_transitions = tf.shape(tag_indices)[1] - 1

  # Truncate by one on each side of the sequence to get the start and end
  # indices of each transition.
  start_tag_indices = tf.slice(tag_indices, [0, 0], [-1, num_transitions])
  end_tag_indices = tf.slice(tag_indices, [0, 1], [-1, num_transitions])

  # Encode the indices in a flattened representation.
  flattened_transition_indices = start_tag_indices * num_tags + end_tag_indices
  flattened_transition_params = tf.reshape(transition_params, [-1])

  # Get the binary scores based on the flattened representation.
  binary_scores = tf.gather(flattened_transition_params,
                            flattened_transition_indices)

  masks = tf.sequence_mask(
      sequence_lengths,
      maxlen=tf.shape(tag_indices)[1],
      dtype=binary_scores.dtype)
  truncated_masks = tf.slice(masks, [0, 1], [-1, -1])
  binary_scores = tf.reduce_sum(binary_scores * truncated_masks, 1)
  return binary_scores
//  [python]
def decode(carbon_record):
  """Parses an image and label from the given `carbon_reader`."""

  image_raw = getattr(carbon_record, 'image')
  image = tf.decode_raw(image_raw, tf.uint8)

  label_raw = getattr(carbon_record, 'digit')
  label = tf.to_int64(label_raw)

  return label, image
//  [python]
def decode(carbon_record):
  """Parses an image and label from the given `carbon_reader`."""

  image_raw = getattr(carbon_record, 'image')
  image = tf.decode_raw(image_raw, tf.uint8)

  label_raw = getattr(carbon_record, 'digit')
  label = tf.cast(label_raw, tf.int64)

  return label, image
//  [python]
def crf_forward(
    inputs,
    state,
    transition_params,
    sequence_lengths,
):
  """Computes the alpha values in a linear-chain CRF.

  See http://www.cs.columbia.edu/~mcollins/fb.pdf for reference.

  Args:
    inputs: A [batch_size, num_tags] matrix of unary potentials.
    state: A [batch_size, num_tags] matrix containing the previous alpha
      values.
    transition_params: A [num_tags, num_tags] matrix of binary potentials.
      This matrix is expanded into a [1, num_tags, num_tags] in preparation
      for the broadcast summation occurring within the cell.
    sequence_lengths: A [batch_size] vector of true sequence lengths.

  Returns:
    new_alphas: A [batch_size, num_tags] matrix containing the
        new alpha values.
  """
  sequence_lengths = tf.to_int32(sequence_lengths)

  last_index = tf.maximum(
      tf.constant(0, dtype=sequence_lengths.dtype), sequence_lengths - 1)
  inputs = tf.transpose(inputs, [1, 0, 2])
  transition_params = tf.expand_dims(transition_params, 0)

  def _scan_fn(scan_state, scan_inputs):
    scan_state = tf.expand_dims(scan_state, 2)
    transition_scores = scan_state + transition_params
    new_alphas = scan_inputs + tf.reduce_logsumexp(transition_scores, [1])
    return new_alphas

  all_alphas = tf.transpose(tf.scan(_scan_fn, inputs, state), [1, 0, 2])
  # add first state for sequences of length 1
  all_alphas = tf.concat([tf.expand_dims(state, 1), all_alphas], 1)

  idxs = tf.stack([tf.range(tf.shape(last_index)[0]), last_index], axis=1)
  return tf.gather_nd(all_alphas, idxs)