jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def crf_sequence_score(
    inputs: TensorLike,
    tag_indices: TensorLike,
    sequence_lengths: TensorLike,
    transition_params: TensorLike,
) -> tf.Tensor:
    """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)

    # If max_seq_len is 1, we skip the score calculation and simply gather the
    # unary potentials of the single tag.
    def _single_seq_fn() -> TensorLike:
        batch_size = tf.shape(inputs, out_type=tf.int32)[0]
        batch_inds = tf.reshape(tf.range(batch_size), [-1, 1])
        indices = tf.concat([batch_inds, tf.zeros_like(batch_inds)], axis=1)

        tag_inds = tf.gather_nd(tag_indices, indices)
        tag_inds = tf.reshape(tag_inds, [-1, 1])
        indices = tf.concat([indices, tag_inds], axis=1)

        sequence_scores = tf.gather_nd(inputs, indices)

        sequence_scores = tf.where(
            tf.less_equal(sequence_lengths, 0),
            tf.zeros_like(sequence_scores),
            sequence_scores,
        )
        return sequence_scores

    def _multi_seq_fn() -> TensorLike:
        # 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 tf.cond(tf.equal(tf.shape(inputs)[1], 1), _single_seq_fn, _multi_seq_fn)
//  [python]
def tf_times():
    """Returns (time since start, time since last) as a tensorflow op."""
    # Keep track of start and last times
    with tf.init_scope():
        init = tf.timestamp()

    def make(name):
        return tf.Variable(init, name=name, trainable=False, use_resource=True)

    start = make('start_time')
    last = make('last_time')

    # Get new time and update last
    now = tf.timestamp()
    prev = last.read_value()
    with tf.control_dependencies([prev]):
        with tf.control_dependencies([last.assign(now)]):
            return tf.to_float(now - start.read_value()), tf.to_float(now - prev)
//  [python]
def tf_times():
    """Returns (time since start, time since last) as a tensorflow op."""
    # Keep track of start and last times
    with tf.init_scope():
        init = tf.timestamp()

    def make(name):
        return tf.Variable(init, name=name, trainable=False, use_resource=True)

    start = make('start_time')
    last = make('last_time')

    # Get new time and update last
    now = tf.timestamp()
    prev = last.read_value()
    with tf.control_dependencies([prev]):
        with tf.control_dependencies([last.assign(now)]):
            return tf.cast(now - start.read_value(), tf.float32), tf.cast(now - prev, tf.float32)
//  [python]
def crf_log_norm(
    inputs: TensorLike, sequence_lengths: TensorLike, transition_params: TensorLike
) -> tf.Tensor:
    """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=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])

    # If max_seq_len is 1, we skip the algorithm and simply reduce_logsumexp
    # over the "initial state" (the unary potentials).
    def _single_seq_fn() -> TensorLike:
        log_norm = tf.reduce_logsumexp(first_input, [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

    def _multi_seq_fn() -> TensorLike:
        """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 tf.cond(tf.equal(tf.shape(inputs)[1], 1), _single_seq_fn, _multi_seq_fn)
//  [python]
def crf_log_norm(
    inputs: TensorLike, sequence_lengths: TensorLike, transition_params: TensorLike
) -> tf.Tensor:
    """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])

    # If max_seq_len is 1, we skip the algorithm and simply reduce_logsumexp
    # over the "initial state" (the unary potentials).
    def _single_seq_fn() -> TensorLike:
        log_norm = tf.reduce_logsumexp(first_input, [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

    def _multi_seq_fn() -> TensorLike:
        """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 tf.cond(tf.equal(tf.shape(inputs)[1], 1), _single_seq_fn, _multi_seq_fn)
//  [python]
def estimate_face_normals(v, f_arr, f_tensor):
    face_normals = tf.reshape(TriNormals(v, f_tensor),(-1, 3)) # this is *unnormalized* face normals (abs != 1), if you need normalized norms, use TriNormals()
    face_normals = t.cast(face_normals, tf.float32)
    return face_normals
//  [python]
def estimate_face_normals(v, f_arr, f_tensor):
    face_normals = tf.reshape(TriNormals(v, f_tensor),(-1, 3)) # this is *unnormalized* face normals (abs != 1), if you need normalized norms, use TriNormals()
    face_normals = tf.cast(face_normals, tf.float32)
    return face_normals
//  [python]
def face_normal_loss_calc(pred, gt, f):
    '''
    This version works as long as the values feed into pred is *not all zeros* (otherwise will produce nan),
    but this situation should never happen in evaluation/test phase, because if a batch is all filled with 0 placeholders,
    this batch will not be created.
    '''
    batch_size = int(gt.get_shape()[0])
    f_tensor = tf.cast(f.astype('int32'), tf.int32)
    normals_pred = tf.stack([estimate_face_normals(pred[x], f, f_tensor) for x in range(batch_size)])
    normals_gt = tf.stack([estimate_face_normals(gt[x], f, f_tensor) for x in range(batch_size)])

    zeros = tf.zeros(shape=tf.shape(normals_gt))

    idx_nonzeros = tf.where(tf.norm(tf.norm(normals_gt - zeros, axis=-1), axis=-1) > 1e-14)

    cos = tf.reduce_sum(tf.multiply(normals_pred, normals_gt), axis=-1)
    cos_abs = tf.abs(cos)
    normal_loss = 1 - cos_abs

    result = tf.gather_nd(normal_loss, idx_nonzeros)
    return tf.reduce_mean(result)
//  [python]
def estimate_vertex_normals(v, f_arr, f_tensor):
    '''
    TF implementation of psbody.mesh.mesh.estimate_vertex_normals
    calculate *scaled* (norm==1) normals of all the vertices v
    args:
        v: vertices
        f_arr: faces, nparray
        f_tensor: same faces as f_arr, but stored in a tensor
    returns:
        array of surface normals, shape of (#verts, 3), already centered and scaled (abs=1),
        i.e. only indicates surface normal directions
    '''
    face_normals = tf.reshape(TriNormals(v, f_tensor),(-1, 3)) # this is *unnormalized* face normals (abs != 1), if you need normalized norms, use TriNormals()
    face_normals = t.cast(face_normals, tf.float32)
    ftov = faces_by_vertex(v, f_arr, as_sparse_matrix=True) # scipy sparse tensor
    ftov = ftov.astype('float32')
    ftov = sparse2tfsparse(ftov)
    non_scaled_normals = tf.sparse_tensor_dense_matmul(ftov, face_normals)
    norms = tf.norm(non_scaled_normals, axis=1)

    indices = tf.equal(norms, 0.0)
    mask = t.cast(indices, norms.dtype) # a mask, 1 where norms==0, 0 otherwise
    norms = tf.add(norms, mask)

    normalized_vert_normals = tf.transpose(tf.divide(tf.transpose(non_scaled_normals), norms))

    return normalized_vert_normals
//  [python]
def estimate_vertex_normals(v, f_arr, f_tensor):
    '''
    TF implementation of psbody.mesh.mesh.estimate_vertex_normals
    calculate *scaled* (norm==1) normals of all the vertices v
    args:
        v: vertices
        f_arr: faces, nparray
        f_tensor: same faces as f_arr, but stored in a tensor
    returns:
        array of surface normals, shape of (#verts, 3), already centered and scaled (abs=1),
        i.e. only indicates surface normal directions
    '''
    face_normals = tf.reshape(TriNormals(v, f_tensor),(-1, 3)) # this is *unnormalized* face normals (abs != 1), if you need normalized norms, use TriNormals()
    face_normals = tf.cast(face_normals, tf.float32)
    ftov = faces_by_vertex(v, f_arr, as_sparse_matrix=True) # scipy sparse tensor
    ftov = ftov.astype('float32')
    ftov = sparse2tfsparse(ftov)
    non_scaled_normals = tf.sparse_tensor_dense_matmul(ftov, face_normals)
    norms = tf.norm(non_scaled_normals, axis=1)

    indices = tf.equal(norms, 0.0)
    mask = tf.cast(indices, norms.dtype) # a mask, 1 where norms==0, 0 otherwise
    norms = tf.add(norms, mask)

    normalized_vert_normals = tf.transpose(tf.divide(tf.transpose(non_scaled_normals), norms))

    return normalized_vert_normals
//  [python]
def call(self, input_features: tf.Tensor) -> tf.Tensor:
        # TF Conv1D assumes Batch x Time x Channels, same as the input
        hidden_states = t.cast(input_features, tf.float32)
        for i, conv in enumerate(self.conv_layers):
            # equivalent to `padding=k // 2` on PT's `nn.Conv1d`
            pad_len = self.kernel_sizes[i] // 2
            hidden_shapes = shape_list(hidden_states)
            hidden_states = tf.concat(
                (
                    tf.zeros((hidden_shapes[0], pad_len, hidden_shapes[2])),
                    hidden_states,
                    tf.zeros((hidden_shapes[0], pad_len, hidden_shapes[2])),
                ),
                axis=1,
            )

            hidden_states = conv(hidden_states)
            hidden_states = glu(hidden_states, axis=2)  # GLU over the Channel dimension
        return hidden_states
//  [python]
def call(self, input_features: tf.Tensor) -> tf.Tensor:
        # TF Conv1D assumes Batch x Time x Channels, same as the input
        hidden_states = tf.cast(input_features, tf.float32)
        for i, conv in enumerate(self.conv_layers):
            # equivalent to `padding=k // 2` on PT's `nn.Conv1d`
            pad_len = self.kernel_sizes[i] // 2
            hidden_shapes = shape_list(hidden_states)
            hidden_states = tf.concat(
                (
                    tf.zeros((hidden_shapes[0], pad_len, hidden_shapes[2])),
                    hidden_states,
                    tf.zeros((hidden_shapes[0], pad_len, hidden_shapes[2])),
                ),
                axis=1,
            )

            hidden_states = conv(hidden_states)
            hidden_states = glu(hidden_states, axis=2)  # GLU over the Channel dimension
        return hidden_states
//  [python]
def shift_tokens_right(input_ids: tf.Tensor, pad_token_id: int, decoder_start_token_id: int):
    pad_token_id = tf.cast(pad_token_id, input_ids.dtype)
    decoder_start_token_id = tf.cast(decoder_start_token_id, input_ids.dtype)
    start_tokens = tf.fill(
        (shape_list(input_ids)[0], 1), tf.convert_to_tensor(decoder_start_token_id, input_ids.dtype)
    )
    shifted_input_ids = tf.concat([start_tokens, input_ids[:, :-1]], -1)
    # replace possible -100 values in labels by `pad_token_id`
    shifted_input_ids = tf.where(
        shifted_input_ids == -100,
        tf.fill(shape_list(shifted_input_ids), tf.convert_to_tensor(pad_token_id, input_ids.dtype)),
        shifted_input_ids,
    )

    # "Verify that `labels` has only positive values and -100"
    assert_gte0 = tf.debugging.assert_greater_equal(shifted_input_ids, tf.constant(0, dtype=input_ids.dtype))

    # Make sure the assertion op is called by wrapping the result in an identity no-op
    with tf.control_dependencies([assert_gte0]):
        shifted_input_ids = tf.identity(shifted_input_ids)

    return shifted_input_ids
//  [python]
def create_position_ids_from_input_ids(
        input_ids: tf.Tensor, padding_idx: int, past_key_values_length: Optional[int] = 0
    ) -> tf.Tensor:
        """
        Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding
        symbols are ignored. This is modified from fairseq's `utils.make_positions`.

        Args:
            x: tf.Tensor x:
        Returns: tf.Tensor
        """
        mask = t.cast(t.math.not_equal(input_ids, padding_idx), dtype=t.int32)
        incremental_indices = (t.math.cumsum(mask, axis=1) + past_key_values_length) * mask
        return t.cast(incremental_indices, dtype=t.int64) + padding_idx
//  [python]
def create_position_ids_from_input_ids(
        input_ids: tf.Tensor, padding_idx: int, past_key_values_length: Optional[int] = 0
    ) -> tf.Tensor:
        """
        Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding
        symbols are ignored. This is modified from fairseq's `utils.make_positions`.

        Args:
            x: tf.Tensor x:
        Returns: tf.Tensor
        """
        mask = tf.cast(tf.math.not_equal(input_ids, padding_idx), dtype=tf.int32)
        incremental_indices = (tf.math.cumsum(mask, axis=1) + past_key_values_length) * mask
        return tf.cast(incremental_indices, dtype=tf.int64) + padding_idx
//  [python]
def process(self, total_probs, beam_width, vocab_size):
        '''
        inputs:
            total_probs: float tensor, [batch_size * beam_width, vocab_size]
        
        outputs:
            sample_ids: int tensor, [batch_size * beam_width]
            beam_ids: int tensor, [batch_size * beam_width]
        
        1. calculate hypothese for each beam
        2. Intra-sibling ordering
        3. rewrite scores
        4. choose top K hypothese
        '''
        
        total_probs = tf.reshape(total_probs, [-1, beam_width, vocab_size]) # [batch size, beam width, vocab size]

        sibling_score = tf.to_float(tf.range(1, beam_width+1)) * self.diversity_rate # [beam_width]
        
        scores, ids = tf.nn.top_k(total_probs, beam_width) # [batch size, beam width, beam width]
        scores = tf.add(scores, sibling_score) # [batch size, beam width, beam width]
        
        scores = tf.reshape(scores, [-1, beam_width * beam_width])
        ids = ids + tf.expand_dims(tf.expand_dims(tf.range(0, beam_width) * vocab_size, 0), -1)
        ids = tf.reshape(ids, [-1, beam_width * beam_width])
        
        _, final_ids = tf.nn.top_k(scores, beam_width) # [batch size, beam width]
        
        batch_size = tf.shape(final_ids)[0]
        final_ids = tf.reshape(final_ids, [-1, 1])
        batch_index = tf.range(0, batch_size) 
        batch_index = tf.reshape(batch_index, [-1, 1])
        batch_index = tf.tile(batch_index, [1, beam_width])
        batch_index = tf.reshape(batch_index, [-1, 1])

        index = tf.concat([batch_index, final_ids ], axis=1)
        sample_ids = tf.gather_nd(ids, index)
         
        return sample_ids
//  [python]
def process(self, total_probs, beam_width, vocab_size):
        '''
        inputs:
            total_probs: float tensor, [batch_size * beam_width, vocab_size]
        
        outputs:
            sample_ids: int tensor, [batch_size * beam_width]
            beam_ids: int tensor, [batch_size * beam_width]
        
        1. calculate hypothese for each beam
        2. Intra-sibling ordering
        3. rewrite scores
        4. choose top K hypothese
        '''
        
        total_probs = tf.reshape(total_probs, [-1, beam_width, vocab_size]) # [batch size, beam width, vocab size]

        sibling_score = tf.cast(tf.range(1, beam_width+1), total_probs.dtype) * self.diversity_rate # [beam_width]
        
        scores, ids = tf.nn.top_k(total_probs, beam_width) # [batch size, beam width, beam width]
        scores = tf.add(scores, sibling_score) # [batch size, beam width, beam width]
        
        scores = tf.reshape(scores, [-1, beam_width * beam_width])
        ids = ids + tf.expand_dims(tf.expand_dims(tf.range(0, beam_width) * vocab_size, 0), -1)
        ids = tf.reshape(ids, [-1, beam_width * beam_width])
        
        _, final_ids = tf.nn.top_k(scores, beam_width) # [batch size, beam width]
        
        batch_size = tf.shape(final_ids)[0]
        final_ids = tf.reshape(final_ids, [-1, 1])
        batch_index = tf.range(0, batch_size) 
        batch_index = tf.reshape(batch_index, [-1, 1])
        batch_index = tf.tile(batch_index, [1, beam_width])
        batch_index = tf.reshape(batch_index, [-1, 1])

        index = tf.concat([batch_index, final_ids ], axis=1)
        sample_ids = tf.gather_nd(ids, index)
         
        return sample_ids
//  [python]
def _get_feature_vector_attention_mask(self, feature_vector_length, attention_mask):
        # generate creates 3D attention mask, because of the shape of input_features
        # convert it to 2D if thats the case
        if len(attention_mask.shape) > 2:
            attention_mask = attention_mask[:, :, -1]

        subsampled_lengths = self._get_feat_extract_output_lengths(tf.math.reduce_sum(attention_mask, -1))
        bsz = shape_list(attention_mask)[0]
        indices = tf.concat(
            (
                tf.expand_dims(tf.range(bsz, dtype=attention_mask.dtype), -1),
                tf.expand_dims(subsampled_lengths - 1, -1),
            ),
            axis=-1,
        )
        attention_mask = tf.scatter_nd(indices=indices, updates=tf.ones(bsz), shape=[bsz, feature_vector_length])
        attention_mask = tf.reverse(tf.math.cumsum(tf.reverse(attention_mask, [-1]), -1), [-1])
        attention_mask = tf.cast(attention_mask, tf.int64)
        return attention_mask
//  [python]
def _get_feature_vector_attention_mask(self, feature_vector_length, attention_mask):
        # generate creates 3D attention mask, because of the shape of input_features
        # convert it to 2D if thats the case
        if len(attention_mask.shape) > 2:
            attention_mask = attention_mask[:, :, -1]

        subsampled_lengths = self._get_feat_extract_output_lengths(tf.math.reduce_sum(attention_mask, -1))
        bsz = shape_list(attention_mask)[0]
        indices = tf.concat(
            (
                tf.expand_dims(tf.range(bsz, dtype=attention_mask.dtype), -1),
                tf.expand_dims(subsampled_lengths - 1, -1),
            ),
            axis=-1,
        )
        attention_mask = tf.scatter_nd(indices=indices, updates=tf.ones(bsz), shape=[bsz, feature_vector_length])
        attention_mask = tf.cast(tf.reverse(tf.math.cumsum(tf.reverse(attention_mask, [-1]), -1), [-1]), tf.int64)
        return attention_mask
//  [python]
def loss(self, rollouts):
        values = rollouts['values']
        old_logprob = rollouts['logprobs']
        rewards = rollouts['rewards']
        with tf.name_scope('ppo_loss'):
            if self.hparams.ppo.whiten_rewards:
                rewards = utils.whiten(rewards, shift_mean=False)

            lastgaelam = 0
            advantages_reversed = []
            gen_length = self.hparams.task.response_length
            for t in reversed(range(gen_length)):
                nextvalues = values[:, t + 1] if t < gen_length - 1 else 0.0
                delta = rewards[:, t] + self.hparams.ppo.gamma * nextvalues - values[:, t]
                lastgaelam = delta + self.hparams.ppo.gamma * self.hparams.ppo.lam * lastgaelam
                advantages_reversed.append(lastgaelam)
            advantages = tf.stack(advantages_reversed[::-1], axis=1)
            returns = advantages + values

            advantages = utils.whiten(advantages)
            advantages = tf.stop_gradient(advantages)  # Shouldn't do anything, but better not to think about it

            outputs = self.policy.analyze_responses_op(rollouts['queries'], rollouts['responses'])

            vpred = outputs['values']
            vpredclipped = tf.clip_by_value(vpred, values - self.hparams.ppo.cliprange_value, values + self.hparams.ppo.cliprange_value)
            vf_losses1 = tf.square(vpred - returns)
            vf_losses2 = tf.square(vpredclipped - returns)
            vf_loss = .5 * tf.reduce_mean(tf.maximum(vf_losses1, vf_losses2))
            vf_clipfrac = tf.reduce_mean(tf.to_float(tf.greater(vf_losses2, vf_losses1)))

            logprob = outputs['logprobs']
            ratio = tf.exp(logprob - old_logprob)
            pg_losses = -advantages * ratio
            pg_losses2 = -advantages * tf.clip_by_value(ratio, 1.0 - self.hparams.ppo.cliprange, 1.0 + self.hparams.ppo.cliprange)
            pg_loss = tf.reduce_mean(tf.maximum(pg_losses, pg_losses2))
            pg_clipfrac = tf.reduce_mean(tf.to_float(tf.greater(pg_losses2, pg_losses)))

            loss = pg_loss + self.hparams.ppo.vf_coef * vf_loss

            entropy = tf.reduce_mean(outputs['entropies'])
            approxkl = .5 * tf.reduce_mean(tf.square(logprob - old_logprob))

            return_mean, return_var = tf.nn.moments(returns, axes=list(range(returns.shape.ndims)))
            value_mean, value_var = tf.nn.moments(values, axes=list(range(values.shape.ndims)))

            stats = dict(
                loss=dict(policy=pg_loss, value=vf_loss, total=loss),
                policy=dict(entropy=entropy, approxkl=approxkl, clipfrac=pg_clipfrac),
                returns=dict(mean=return_mean, var=return_var),
                val=dict(vpred=tf.reduce_mean(vpred), error=tf.reduce_mean((vpred - returns) ** 2),
                         clipfrac=vf_clipfrac, mean=value_mean, var=value_var)
            )
            return loss, utils.flatten_dict(stats, sep='/')
//  [python]
def loss(self, rollouts):
        values = rollouts['values']
        old_logprob = rollouts['logprobs']
        rewards = rollouts['rewards']
        with tf.name_scope('ppo_loss'):
            if self.hparams.ppo.whiten_rewards:
                rewards = utils.whiten(rewards, shift_mean=False)

            lastgaelam = 0
            advantages_reversed = []
            gen_length = self.hparams.task.response_length
            for t in reversed(range(gen_length)):
                nextvalues = values[:, t + 1] if t < gen_length - 1 else 0.0
                delta = rewards[:, t] + self.hparams.ppo.gamma * nextvalues - values[:, t]
                lastgaelam = delta + self.hparams.ppo.gamma * self.hparams.ppo.lam * lastgaelam
                advantages_reversed.append(lastgaelam)
            advantages = tf.stack(advantages_reversed[::-1], axis=1)
            returns = advantages + values

            advantages = utils.whiten(advantages)
            advantages = tf.stop_gradient(advantages)  # Shouldn't do anything, but better not to think about it

            outputs = self.policy.analyze_responses_op(rollouts['queries'], rollouts['responses'])

            vpred = outputs['values']
            vpredclipped = tf.clip_by_value(vpred, values - self.hparams.ppo.cliprange_value, values + self.hparams.ppo.cliprange_value)
            vf_losses1 = tf.square(vpred - returns)
            vf_losses2 = tf.square(vpredclipped - returns)
            vf_loss = .5 * tf.reduce_mean(tf.maximum(vf_losses1, vf_losses2))
            vf_clipfrac = tf.reduce_mean(tf.cast(tf.greater(vf_losses2, vf_losses1), tf.float32))

            logprob = outputs['logprobs']
            ratio = tf.exp(logprob - old_logprob)
            pg_losses = -advantages * ratio
            pg_losses2 = -advantages * tf.clip_by_value(ratio, 1.0 - self.hparams.ppo.cliprange, 1.0 + self.hparams.ppo.cliprange)
            pg_loss = tf.reduce_mean(tf.maximum(pg_losses, pg_losses2))
            pg_clipfrac = tf.reduce_mean(tf.cast(tf.greater(pg_losses2, pg_losses), tf.float32))

            loss = pg_loss + self.hparams.ppo.vf_coef * vf_loss

            entropy = tf.reduce_mean(outputs['entropies'])
            approxkl = .5 * tf.reduce_mean(tf.square(logprob - old_logprob))

            return_mean, return_var = tf.nn.moments(returns, axes=list(range(returns.shape.ndims)))
            value_mean, value_var = tf.nn.moments(values, axes=list(range(values.shape.ndims)))

            stats = dict(
                loss=dict(policy=pg_loss, value=vf_loss, total=loss),
                policy=dict(entropy=entropy, approxkl=approxkl, clipfrac=pg_clipfrac),
                returns=dict(mean=return_mean, var=return_var),
                val=dict(vpred=tf.reduce_mean(vpred), error=tf.reduce_mean((vpred - returns) ** 2),
                         clipfrac=vf_clipfrac, mean=value_mean, var=value_var)
            )
            return loss, utils.flatten_dict(stats, sep='/')
//  [python]
def loss_function(self, pred, target):
        """
        L2 loss function calculated with prediction and target.
        :param pred: The predicted image.
        :param target: The target image.
        :return: L2 loss of (pred - target) / batch_size
        """
        return tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.cast(target, tf.float32), logits=tf.cast(pred, tf.float32)))
//  [python]
def forward(self, layer_in):
        input_size = tf.shape(layer_in)
        patches = tf.extract_image_patches(images=layer_in, ksizes=[1, self.kernel_size, self.kernel_size, 1],
                                           strides=[1, self.strides, self.strides, 1], rates=[1, 1, 1, 1],
                                           padding="SAME")
        patches_flatten = tf.reshape(patches,
                                     [input_size[0], -1, self.in_channel * self.kernel_size * self.kernel_size])
        patches_infpad = tf.where(tf.less(patches_flatten, 0.1),
                                  self.MAX_SPIKE_TIME * tf.ones_like(patches_flatten), patches_flatten)
        img_raw = tf.map_fn(self.kernel.forward, patches_infpad)
        img_reshaped = tf.reshape(img_raw,
                                  [input_size[0], tf.to_int32(tf.math.ceil(input_size[1] / self.strides)),
                                   tf.to_int32(tf.math.ceil(input_size[2] / self.strides)),
                                   self.out_channel])
        return img_reshaped
//  [python]
def forward(self, layer_in):
        input_size = tf.shape(layer_in)
        patches = tf.extract_image_patches(images=layer_in, ksizes=[1, self.kernel_size, self.kernel_size, 1],
                                           strides=[1, self.strides, self.strides, 1], rates=[1, 1, 1, 1],
                                           padding="SAME")
        patches_flatten = tf.reshape(patches,
                                     [input_size[0], -1, self.in_channel * self.kernel_size * self.kernel_size])
        patches_infpad = tf.where(tf.less(patches_flatten, 0.1),
                                  self.MAX_SPIKE_TIME * tf.ones_like(patches_flatten), patches_flatten)
        img_raw = tf.map_fn(self.kernel.forward, patches_infpad)
        img_reshaped = tf.reshape(img_raw,
                                  [input_size[0], tf.cast(tf.math.ceil(input_size[1] / self.strides), tf.int32),
                                   tf.cast(tf.math.ceil(input_size[2] / self.strides), tf.int32),
                                   self.out_channel])
        return img_reshaped
//  [python]
def loss_func(both):
    """
    function to calculate loss, refer to paper p.7, formula 11
    :param both: a tensor, it put both layer output and expected output together, its' shape
            is [batch_size,out_size*2], where the left part is layer output(real output), right part is
            the label of input(expected output), the tensor both should be looked like this:
            [[2.13,3.56,7.33,3.97,...0,0,1,0,...]
             [3.14,5.56,2.54,15.6,...0,0,0,1,...]...]
                ↑                   ↑
             layer output           label of input
    :return: a tensor, it is a scalar of loss
    """
    output = tf.slice(both, [0], [tf.to_int32(tf.shape(both)[0] / 2)])
    index = tf.slice(both, [tf.to_int32(tf.shape(both)[0] / 2)],
                     [tf.to_int32(tf.shape(both)[0] / 2)])
    z1 = tf.exp(tf.subtract(0., tf.reduce_sum(tf.multiply(output, index))))
    z2 = tf.reduce_sum(tf.exp(tf.subtract(0., output)))
    loss = tf.subtract(
        0., tf.log(
            tf.clip_by_value(tf.divide(
                z1, tf.clip_by_value(
                    z2, 1e-10, 1e10)), 1e-10, 1)))
    return loss
//  [python]
def loss_func(both):
    """
    function to calculate loss, refer to paper p.7, formula 11
    :param both: a tensor, it put both layer output and expected output together, its' shape
            is [batch_size,out_size*2], where the left part is layer output(real output), right part is
            the label of input(expected output), the tensor both should be looked like this:
            [[2.13,3.56,7.33,3.97,...0,0,1,0,...]
             [3.14,5.56,2.54,15.6,...0,0,0,1,...]...]
                ↑                   ↑
             layer output           label of input
    :return: a tensor, it is a scalar of loss
    """
    output = tf.slice(both, [0], [tf.cast(tf.shape(both)[0] / 2, tf.int32)])
    index = tf.slice(both, [tf.cast(tf.shape(both)[0] / 2, tf.int32)],
                     [tf.cast(tf.shape(both)[0] / 2, tf.int32)])
    z1 = tf.exp(tf.subtract(0., tf.reduce_sum(tf.multiply(output, index))))
    z2 = tf.reduce_sum(tf.exp(tf.subtract(0., output)))
    loss = tf.subtract(
        0., tf.log(
            tf.clip_by_value(tf.divide(
                z1, tf.clip_by_value(
                    z2, 1e-10, 1e10)), 1e-10, 1)))
    return loss
//  [python]
def _call_diag_cov(self, x):
        pooled, self.idx = self._pool2d_argmax(self.layer.input, 
                                          pool_size=self.layer.pool_size,
                                          strides=self.layer.strides,
                                          padding=self.layer.padding)
        shape_x = tf.shape(x)
        shape_pooled = tf.shape(pooled)
        shape_arg_max = tf.shape(self.idx)
        x_flat = tf.reshape(x, [shape_x[0], shape_x[1]*shape_x[2]*shape_x[3]])
        arg_max_flat = tf.reshape(self.idx, [shape_arg_max[0], shape_arg_max[1]*shape_arg_max[2]*shape_arg_max[3]])
        arg_max_flat = tf.to_int32(arg_max_flat)
        out = tf.batch_gather(x_flat, arg_max_flat)
        out = tf.reshape(out, [shape_pooled[0], shape_pooled[1], shape_pooled[2], shape_pooled[3]])
        return out
//  [python]
def _call_diag_cov(self, x):
        pooled, self.idx = self._pool2d_argmax(self.layer.input, 
                                          pool_size=self.layer.pool_size,
                                          strides=self.layer.strides,
                                          padding=self.layer.padding)
        shape_x = tf.shape(x)
        shape_pooled = tf.shape(pooled)
        shape_arg_max = tf.shape(self.idx)
        x_flat = tf.reshape(x, [shape_x[0], shape_x[1]*shape_x[2]*shape_x[3]])
        arg_max_flat = tf.reshape(self.idx, [shape_arg_max[0], shape_arg_max[1]*shape_arg_max[2]*shape_arg_max[3]])
        arg_max_flat = tf.cast(arg_max_flat, tf.int32)
        out = tf.batch_gather(x_flat, arg_max_flat)
        out = tf.reshape(out, [shape_pooled[0], shape_pooled[1], shape_pooled[2], shape_pooled[3]])
        return out
//  [python]
def _train(self, experience, weights):
    experience = self._as_trajectory(experience)

    with tf.GradientTape() as tape:
      loss_info = self._loss(experience, weights=weights, training=True)

    variables_to_train = self._variables_to_train()
    if not variables_to_train:
      logging.info('No variable to train in the agent.')
      return loss_info

    grads = tape.gradient(loss_info.loss, variables_to_train)
    # Tuple is used for py3, where zip is a generator producing values once.
    grads_and_vars = tuple(zip(grads, variables_to_train))
    if self._gradient_clipping is not None:
      grads_and_vars = eager_utils.clip_gradient_norms(
          grads_and_vars, self._gradient_clipping
      )

    if self._summarize_grads_and_vars:
      eager_utils.add_variables_summaries(
          grads_and_vars, self.train_step_counter
      )
      eager_utils.add_gradients_summaries(
          grads_and_vars, self.train_step_counter
      )

    self._optimizer.apply_gradients(grads_and_vars)
    self.train_step_counter.assign_add(1)
    if self._num_samples_list:
      actions_flattened = tf.reshape(experience.action, [-1])
      if self._accepts_per_arm_features:
        num_samples_per_action_current = [
            tf.cast(tf.shape(actions_flattened)[0], dtype=tf.int64)
        ]
      else:
        # Compute the number of samples for each action in the current batch.
        num_samples_per_action_current = [
            tf.reduce_sum(tf.cast(tf.equal(actions_flattened, k), tf.int64))
            for k in range(self._num_actions)
        ]
      # Update the number of samples for each action.
      for a, b in zip(self._num_samples_list, num_samples_per_action_current):
        tf.compat.v1.assign_add(a, b)

    return loss_info
//  [python]
def _compute_threshold(x):
    """
    ref: https://github.com/XJTUWYD/TWN
    Computing the threshold.
    """
    x_sum = tf.reduce_sum(input_tensor=tf.abs(x), axis=None, keepdims=False, name=None)
    threshold = tf.compat.v1.div(x_sum, tf.compat.v1.cast(tf.size(input=x), tf.float32), name=None)
    threshold = tf.multiply(0.7, threshold, name=None)
    return threshold
//  [python]
def _compute_threshold(x):
    """
    ref: https://github.com/XJTUWYD/TWN
    Computing the threshold.
    """
    x_sum = tf.reduce_sum(input_tensor=tf.abs(x), axis=None, keepdims=False, name=None)
    # threshold = tf.compat.v1.div(x_sum, tf.cast(tf.size(input=x), tf.float32), name=None)
    threshold = tf.math.divide(x_sum, tf.cast(tf.size(input=x), tf.float32), name=None)
    threshold = tf.multiply(0.7, threshold, name=None)
    return threshold
//  [python]
def _points_to_voxel_indices(points, grid_cell_size):
  """Converts points into corresponding voxel indices.

  Maps each point into a voxel grid with cell size given by grid_cell_size.
  For each voxel, it computes a x, y, z index. Also converts the x, y, z index
  to a single number index where there is a one-on-one mapping between
  each x, y, z index value and its corresponding single number index value.

  Args:
    points: A tf.float32 tensor of size [N, 3].
    grid_cell_size: The size of the grid cells in x, y, z dimensions in the
      voxel grid. It should be either a tf.float32 tensor, a numpy array or a
      list of size [3].

  Returns:
    voxel_xyz_indices: A tf.int32 tensor of size [N, 3] containing the x, y, z
      index of the voxel corresponding to each given point.
    voxel_single_number_indices: A tf.int32 tensor of size [N] containing the
      single number index of the voxel corresponding to each given point.
    voxel_start_location: A tf.float32 tensor of size [3] containing the start
      location of the voxels.
  """
  voxel_start_location = tf.reduce_min(points, axis=0)
  voxel_xyz_indices = t.cast(
      tf.math.floordiv(points - voxel_start_location, grid_cell_size),
      dtype=tf.int32)
  voxel_xyz_indices, voxel_single_number_indices = compute_pooled_voxel_indices(
      voxel_xyz_indices=voxel_xyz_indices, pooling_size=(1, 1, 1))
  return voxel_xyz_indices, voxel_single_number_indices, voxel_start_location
//  [python]
def _points_to_voxel_indices(points, grid_cell_size):
  """Converts points into corresponding voxel indices.

  Maps each point into a voxel grid with cell size given by grid_cell_size.
  For each voxel, it computes a x, y, z index. Also converts the x, y, z index
  to a single number index where there is a one-on-one mapping between
  each x, y, z index value and its corresponding single number index value.

  Args:
    points: A tf.float32 tensor of size [N, 3].
    grid_cell_size: The size of the grid cells in x, y, z dimensions in the
      voxel grid. It should be either a tf.float32 tensor, a numpy array or a
      list of size [3].

  Returns:
    voxel_xyz_indices: A tf.int32 tensor of size [N, 3] containing the x, y, z
      index of the voxel corresponding to each given point.
    voxel_single_number_indices: A tf.int32 tensor of size [N] containing the
      single number index of the voxel corresponding to each given point.
    voxel_start_location: A tf.float32 tensor of size [3] containing the start
      location of the voxels.
  """
  voxel_start_location = tf.reduce_min(points, axis=0)
  voxel_xyz_indices = tf.cast(
      tf.math.floordiv(points - voxel_start_location, grid_cell_size),
      dtype=tf.int32)
  voxel_xyz_indices, voxel_single_number_indices = compute_pooled_voxel_indices(
      voxel_xyz_indices=voxel_xyz_indices, pooling_size=(1, 1, 1))
  return voxel_xyz_indices, voxel_single_number_indices, voxel_start_location
//  [python]
def __call__(self, image_size):
    image_height = t.cast(image_size[0], t.float32)
    image_width = t.cast(image_size[1], t.float32)

    k = len(self.scales) * len(self.aspect_ratios)
    aspect_ratios_sqrt = t.cast(t.sqrt(self.aspect_ratios), dtype=t.float32)
    anchor_size = t.cast(self.anchor_size, t.float32)

    # [K]
    anchor_heights = []
    anchor_widths = []
    for scale in self.scales:
      anchor_size_t = anchor_size * scale
      anchor_height = anchor_size_t / aspect_ratios_sqrt
      anchor_width = anchor_size_t * aspect_ratios_sqrt
      anchor_heights.append(anchor_height)
      anchor_widths.append(anchor_width)
    anchor_heights = t.concat(anchor_heights, axis=0)
    anchor_widths = t.concat(anchor_widths, axis=0)
    half_anchor_heights = t.reshape(0.5 * anchor_heights, [1, 1, k])
    half_anchor_widths = t.reshape(0.5 * anchor_widths, [1, 1, k])

    stride = t.cast(self.stride, t.float32)
    # [W]
    cx = t.range(0.5 * stride, image_width + 0.5 * stride, stride)
    # [H]
    cy = t.range(0.5 * stride, image_height + 0.5 * stride, stride)
    # [H, W]
    cx_grid, cy_grid = t.meshgrid(cx, cy)
    # [H, W, 1]
    cx_grid = t.expand_dims(cx_grid, axis=-1)
    cy_grid = t.expand_dims(cy_grid, axis=-1)

    # [H, W, K, 1]
    y_min = t.expand_dims(cy_grid - half_anchor_heights, axis=-1)
    y_max = t.expand_dims(cy_grid + half_anchor_heights, axis=-1)
    x_min = t.expand_dims(cx_grid - half_anchor_widths, axis=-1)
    x_max = t.expand_dims(cx_grid + half_anchor_widths, axis=-1)

    if self.clip_boxes:
      y_min = t.maximum(t.minimum(y_min, image_height), 0.)
      y_max = t.maximum(t.minimum(y_max, image_height), 0.)
      x_min = t.maximum(t.minimum(x_min, image_width), 0.)
      x_max = t.maximum(t.minimum(x_max, image_width), 0.)

    # [H, W, K, 4]
    result = t.concat([y_min, x_min, y_max, x_max], axis=-1)
    shape = result.shape.as_list()
    # [H, W, K * 4]
    return t.reshape(result, [shape[0], shape[1], shape[2] * shape[3]])
//  [python]
def __call__(self, image_size):
    image_height = tf.cast(image_size[0], tf.float32)
    image_width = tf.cast(image_size[1], tf.float32)

    k = len(self.scales) * len(self.aspect_ratios)
    aspect_ratios_sqrt = tf.cast(tf.sqrt(self.aspect_ratios), dtype=tf.float32)
    anchor_size = tf.cast(self.anchor_size, tf.float32)

    # [K]
    anchor_heights = []
    anchor_widths = []
    for scale in self.scales:
      anchor_size_t = anchor_size * scale
      anchor_height = anchor_size_t / aspect_ratios_sqrt
      anchor_width = anchor_size_t * aspect_ratios_sqrt
      anchor_heights.append(anchor_height)
      anchor_widths.append(anchor_width)
    anchor_heights = tf.concat(anchor_heights, axis=0)
    anchor_widths = tf.concat(anchor_widths, axis=0)
    half_anchor_heights = tf.reshape(0.5 * anchor_heights, [1, 1, k])
    half_anchor_widths = tf.reshape(0.5 * anchor_widths, [1, 1, k])

    stride = tf.cast(self.stride, tf.float32)
    # [W]
    cx = tf.range(0.5 * stride, image_width + 0.5 * stride, stride)
    # [H]
    cy = tf.range(0.5 * stride, image_height + 0.5 * stride, stride)
    # [H, W]
    cx_grid, cy_grid = tf.meshgrid(cx, cy)
    # [H, W, 1]
    cx_grid = tf.expand_dims(cx_grid, axis=-1)
    cy_grid = tf.expand_dims(cy_grid, axis=-1)

    # [H, W, K, 1]
    y_min = tf.expand_dims(cy_grid - half_anchor_heights, axis=-1)
    y_max = tf.expand_dims(cy_grid + half_anchor_heights, axis=-1)
    x_min = tf.expand_dims(cx_grid - half_anchor_widths, axis=-1)
    x_max = tf.expand_dims(cx_grid + half_anchor_widths, axis=-1)

    if self.clip_boxes:
      y_min = tf.maximum(tf.minimum(y_min, image_height), 0.)
      y_max = tf.maximum(tf.minimum(y_max, image_height), 0.)
      x_min = tf.maximum(tf.minimum(x_min, image_width), 0.)
      x_max = tf.maximum(tf.minimum(x_max, image_width), 0.)

    # [H, W, K, 4]
    result = tf.concat([y_min, x_min, y_max, x_max], axis=-1)
    shape = result.shape.as_list()
    # [H, W, K * 4]
    return tf.reshape(result, [shape[0], shape[1], shape[2] * shape[3]])
//  [python]
def _pad_or_clip_voxels(voxel_features, voxel_indices, num_valid_voxels,
                        segment_ids, voxels_pad_or_clip_size):
  """Pads or clips voxels."""
  if voxels_pad_or_clip_size:
    num_valid_voxels = tf.minimum(num_valid_voxels, voxels_pad_or_clip_size)
    num_channels = voxel_features.get_shape().as_list()[-1]
    if len(voxel_features.shape.as_list()) == 2:
      output_shape = [voxels_pad_or_clip_size, num_channels]
    elif len(voxel_features.shape.as_list()) == 3:
      num_samples_per_voxel = voxel_features.get_shape().as_list()[1]
      if num_samples_per_voxel is None:
        num_samples_per_voxel = tf.shape(voxel_features)[1]
      output_shape = [
          voxels_pad_or_clip_size, num_samples_per_voxel, num_channels
      ]
    else:
      raise ValueError('voxel_features should be either rank 2 or 3.')
    voxel_features = shape_utils.pad_or_clip_nd(
        tensor=voxel_features, output_shape=output_shape)
    voxel_indices = shape_utils.pad_or_clip_nd(
        tensor=voxel_indices, output_shape=[voxels_pad_or_clip_size, 3])
    valid_segment_ids_mask = tf.to_int32(
        tf.less(segment_ids, num_valid_voxels))
    segment_ids *= valid_segment_ids_mask
  return voxel_features, voxel_indices, num_valid_voxels, segment_ids
//  [python]
def _pad_or_clip_voxels(voxel_features, voxel_indices, num_valid_voxels,
                        segment_ids, voxels_pad_or_clip_size):
  """Pads or clips voxels."""
  if voxels_pad_or_clip_size:
    num_valid_voxels = tf.minimum(num_valid_voxels, voxels_pad_or_clip_size)
    num_channels = voxel_features.get_shape().as_list()[-1]
    if len(voxel_features.shape.as_list()) == 2:
      output_shape = [voxels_pad_or_clip_size, num_channels]
    elif len(voxel_features.shape.as_list()) == 3:
      num_samples_per_voxel = voxel_features.get_shape().as_list()[1]
      if num_samples_per_voxel is None:
        num_samples_per_voxel = tf.shape(voxel_features)[1]
      output_shape = [
          voxels_pad_or_clip_size, num_samples_per_voxel, num_channels
      ]
    else:
      raise ValueError('voxel_features should be either rank 2 or 3.')
    voxel_features = shape_utils.pad_or_clip_nd(
        tensor=voxel_features, output_shape=output_shape)
    voxel_indices = shape_utils.pad_or_clip_nd(
        tensor=voxel_indices, output_shape=[voxels_pad_or_clip_size, 3])
    valid_segment_ids_mask = tf.cast(
        tf.less(segment_ids, num_valid_voxels), dtype=tf.int32)
    segment_ids *= valid_segment_ids_mask
  return voxel_features, voxel_indices, num_valid_voxels, segment_ids
//  [python]
def classification_loss(logits, labels, weights):
  """Computes cross entropy loss between logits and labels.

  Args:
    logits: 2-D [timesteps*batch_size, m] float tensor, where m=1 if
      num_classes=2, otherwise m=num_classes.
    labels: 1-D [timesteps*batch_size] integer tensor.
    weights: 1-D [timesteps*batch_size] float tensor.

  Returns:
    Loss scalar of type float.
  """
  inner_dim = logits.get_shape().as_list()[-1]
  with tf.name_scope('classifier_loss'):
    # Logistic loss
    if inner_dim == 1:
      loss = tf.nn.sigmoid_cross_entropy_with_logits(
          logits=tf.squeeze(logits, -1), labels=tf.to_float(labels))
    # Softmax loss
    else:
      loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
          logits=logits, labels=labels)

    num_lab = _num_labels(weights)
    tf.summary.scalar('num_labels', num_lab)
    return tf.identity(
        tf.reduce_sum(weights * loss) / num_lab, name='classification_xentropy')
//  [python]
def classification_loss(logits, labels, weights):
  """Computes cross entropy loss between logits and labels.

  Args:
    logits: 2-D [timesteps*batch_size, m] float tensor, where m=1 if
      num_classes=2, otherwise m=num_classes.
    labels: 1-D [timesteps*batch_size] integer tensor.
    weights: 1-D [timesteps*batch_size] float tensor.

  Returns:
    Loss scalar of type float.
  """
  inner_dim = logits.get_shape().as_list()[-1]
  with tf.name_scope('classifier_loss'):
    # Logistic loss
    if inner_dim == 1:
      loss = tf.nn.sigmoid_cross_entropy_with_logits(
          logits=tf.squeeze(logits, -1), labels=tf.cast(labels, tf.float32))
    # Softmax loss
    else:
      loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
          logits=logits, labels=labels)

    num_lab = _num_labels(weights)
    tf.summary.scalar('num_labels', num_lab)
    return tf.identity(
        tf.reduce_sum(weights * loss) / num_lab, name='classification_xentropy')