jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _get_output( rnn_logits, sequence_length, lexicon, lexicon_prior ):
    """Create output ops for validation (testing) and prediction
       prediction: Result of CTC beam search decoding
       seq_prob: Score of prediction
    """
    with tf.name_scope("test"):
        if lexicon:
            ctc_blank = (rnn_logits.shape[2]-1)
            lex_prediction,lex_seq_prob = _get_lexicon_output(rnn_logits,
                                                      sequence_length, lexicon )
            if lexicon_prior != None:
                # Need to run both open and closed vocabulary modes
                open_prediction, open_seq_prob = _get_open_output(
                    rnn_logits, sequence_length)
                # Convert top open output prediction to dense values
                # NOTE: What to do if the sparse result is shorter than T?
                # Reshape sparse version of open_prediction?
                open_prediction = tf.to_int32(
                    tf.sparse.to_dense(
                        tf.sparse.reset_shape(
                            open_prediction[0],
                            new_shape=tf.shape(lex_prediction) ),
                        default_value=ctc_blank))
                prediction, seq_prob = _get_merged_output(
                    lex_prediction, lex_seq_prob,
                    open_prediction, open_seq_prob, lexicon_prior )
            else:
                prediction = lex_prediction
                seq_prob = lex_seq_prob
                
            # Match tf.nn.ctc_beam_search_decoder outputs: list of sparse

            # (1) CTCWordBeamSearch returns a dense tensor matching input 
            # sequence length (padded with ctc blanks).
            # We convert to sparse tightly so trailing blanks are trimmed from 
            # the dense_shape of the resulting SparseTensor
            prediction = utils.dense_to_sparse_tight(
                prediction,
                eos_token=ctc_blank )
            # (2) CTCWordBeamSearch returns only top match, so convert to list
            prediction = [prediction]
        else:
            prediction, seq_prob = _get_open_output(rnn_logits, sequence_length)
            
    return prediction, seq_prob
//  [python]
def _get_output( rnn_logits, sequence_length, lexicon, lexicon_prior ):
    """Create output ops for validation (testing) and prediction
       prediction: Result of CTC beam search decoding
       seq_prob: Score of prediction
    """
    with tf.name_scope("test"):
        if lexicon:
            ctc_blank = (rnn_logits.shape[2]-1)
            lex_prediction,lex_seq_prob = _get_lexicon_output(rnn_logits,
                                                      sequence_length, lexicon )
            if lexicon_prior != None:
                # Need to run both open and closed vocabulary modes
                open_prediction, open_seq_prob = _get_open_output(
                    rnn_logits, sequence_length)
                # Convert top open output prediction to dense values
                # NOTE: What to do if the sparse result is shorter than T?
                # Reshape sparse version of open_prediction?
                open_prediction = tf.cast(
                    tf.sparse.to_dense(
                        tf.sparse.reset_shape(
                            open_prediction[0],
                            new_shape=tf.shape(lex_prediction) ),
                        default_value=ctc_blank),
                    tf.int32)
                prediction, seq_prob = _get_merged_output(
                    lex_prediction, lex_seq_prob,
                    open_prediction, open_seq_prob, lexicon_prior )
            else:
                prediction = lex_prediction
                seq_prob = lex_seq_prob
                
            # Match tf.nn.ctc_beam_search_decoder outputs: list of sparse

            # (1) CTCWordBeamSearch returns a dense tensor matching input 
            # sequence length (padded with ctc blanks).
            # We convert to sparse tightly so trailing blanks are trimmed from 
            # the dense_shape of the resulting SparseTensor
            prediction = utils.dense_to_sparse_tight(
                prediction,
                eos_token=ctc_blank )
            # (2) CTCWordBeamSearch returns only top match, so convert to list
            prediction = [prediction]
        else:
            prediction, seq_prob = _get_open_output(rnn_logits, sequence_length)
            
    return prediction, seq_prob
//  [python]
def get_contrast_batch_labels(y):
    '''
    Make contrast labels by taking all the pairwise in y
    y: tensor with shape: (batch_size, )
    returns:   
        tensor with shape: (batch_size * (batch_size-1) // 2, )
    '''
    y_col_vec = tf.reshape(t.cast(y, tf.float32), [-1, 1])
    D_y = pdist_euclidean(y_col_vec)
    d_y = square_to_vec(D_y)
    y_contrasts = t.cast(d_y == 0, tf.int32)
    return y_contrasts
//  [python]
def get_contrast_batch_labels(y):
    '''
    Make contrast labels by taking all the pairwise in y
    y: tensor with shape: (batch_size, )
    returns:   
        tensor with shape: (batch_size * (batch_size-1) // 2, )
    '''
    y_col_vec = tf.reshape(tf.cast(y, tf.float32), [-1, 1])
    D_y = pdist_euclidean(y_col_vec)
    d_y = square_to_vec(D_y)
    y_contrasts = tf.cast(d_y == 0, tf.int32)
    return y_contrasts
//  [python]
def decode(conv_output, anchors, stride, num_class):
    conv_shape       = tf.shape(conv_output)
    batch_size       = conv_shape[0]
    output_size      = conv_shape[1]
    anchor_per_scale = len(anchors)
    conv_output = tf.reshape(conv_output, (batch_size, output_size, output_size, anchor_per_scale, 5 + num_class))
    conv_raw_dxdy = conv_output[:, :, :, :, 0:2]
    conv_raw_dwdh = conv_output[:, :, :, :, 2:4]
    conv_raw_conf = conv_output[:, :, :, :, 4:5]
    conv_raw_prob = conv_output[:, :, :, :, 5: ]
    y = tf.tile(tf.range(output_size, dtype=tf.int32)[:, tf.newaxis], [1, output_size])
    x = tf.tile(tf.range(output_size, dtype=tf.int32)[tf.newaxis, :], [output_size, 1])
    xy_grid = tf.concat([x[:, :, tf.newaxis], y[:, :, tf.newaxis]], axis=-1)
    xy_grid = tf.tile(xy_grid[tf.newaxis, :, :, tf.newaxis, :], [batch_size, 1, 1, anchor_per_scale, 1])
    xy_grid = tf.to_float(xy_grid)
    pred_xy = (tf.sigmoid(conv_raw_dxdy) + xy_grid) * stride
    pred_wh = (tf.exp(conv_raw_dwdh) * anchors) * stride
    pred_xywh = tf.concat([pred_xy, pred_wh], axis=-1)
    pred_conf = tf.sigmoid(conv_raw_conf)
    pred_prob = tf.sigmoid(conv_raw_prob)
    return tf.concat([pred_xywh, pred_conf, pred_prob], axis=-1)
//  [python]
def decode(conv_output, anchors, stride, num_class):
    conv_shape       = tf.shape(conv_output)
    batch_size       = conv_shape[0]
    output_size      = conv_shape[1]
    anchor_per_scale = len(anchors)
    conv_output = tf.reshape(conv_output, (batch_size, output_size, output_size, anchor_per_scale, 5 + num_class))
    conv_raw_dxdy = conv_output[:, :, :, :, 0:2]
    conv_raw_dwdh = conv_output[:, :, :, :, 2:4]
    conv_raw_conf = conv_output[:, :, :, :, 4:5]
    conv_raw_prob = conv_output[:, :, :, :, 5: ]
    y = tf.tile(tf.range(output_size, dtype=tf.int32)[:, tf.newaxis], [1, output_size])
    x = tf.tile(tf.range(output_size, dtype=tf.int32)[tf.newaxis, :], [output_size, 1])
    xy_grid = tf.concat([x[:, :, tf.newaxis], y[:, :, tf.newaxis]], axis=-1)
    xy_grid = tf.tile(xy_grid[tf.newaxis, :, :, tf.newaxis, :], [batch_size, 1, 1, anchor_per_scale, 1])
    xy_grid = tf.cast(xy_grid, tf.float32)
    pred_xy = (tf.sigmoid(conv_raw_dxdy) + xy_grid) * stride
    pred_wh = (tf.exp(conv_raw_dwdh) * anchors) * stride
    pred_xywh = tf.concat([pred_xy, pred_wh], axis=-1)
    pred_conf = tf.sigmoid(conv_raw_conf)
    pred_prob = tf.sigmoid(conv_raw_prob)
    return tf.concat([pred_xywh, pred_conf, pred_prob], axis=-1)
//  [python]
def test_batch_left_padding(self):
        # Confirms that left-padding is working properly
        model, tokenizer, sentences, expected_output_sentences = self._get_beam_search_test_objects()

        inputs = tokenizer(sentences, return_tensors="tf", padding=True)
        inputs_non_padded = tokenizer(sentences[0], return_tensors="tf")
        output_non_padded = model.generate(**inputs_non_padded, do_sample=False, num_beams=2)
        num_paddings = (
            shape_list(inputs_non_padded["input_ids"])[-1]
            - tf.reduce_sum(tf.to_int64(inputs["attention_mask"][-1])).numpy()
        )
        inputs_padded = tokenizer(sentences[1], return_tensors="tf")
        output_padded = model.generate(
            **inputs_padded, do_sample=False, num_beams=2, max_length=model.config.max_length - num_paddings
        )
        non_padded_sentence = tokenizer.decode(output_non_padded[0], skip_special_tokens=True)
        padded_sentence = tokenizer.decode(output_padded[0], skip_special_tokens=True)
        self.assertListEqual(expected_output_sentences, [non_padded_sentence, padded_sentence])
//  [python]
def test_batch_left_padding(self):
        # Confirms that left-padding is working properly
        model, tokenizer, sentences, expected_output_sentences = self._get_beam_search_test_objects()

        inputs = tokenizer(sentences, return_tensors="tf", padding=True)
        inputs_non_padded = tokenizer(sentences[0], return_tensors="tf")
        output_non_padded = model.generate(**inputs_non_padded, do_sample=False, num_beams=2)
        num_paddings = (
            shape_list(inputs_non_padded["input_ids"])[-1]
            - tf.reduce_sum(tf.cast(inputs["attention_mask"][-1], tf.int64)).numpy()
        )
        inputs_padded = tokenizer(sentences[1], return_tensors="tf")
        output_padded = model.generate(
            **inputs_padded, do_sample=False, num_beams=2, max_length=model.config.max_length - num_paddings
        )
        non_padded_sentence = tokenizer.decode(output_non_padded[0], skip_special_tokens=True)
        padded_sentence = tokenizer.decode(output_padded[0], skip_special_tokens=True)
        self.assertListEqual(expected_output_sentences, [non_padded_sentence, padded_sentence])
//  [python]
def get_permutation_loss(labels, predicted_scores):
    tmp_labels = t.cast(t.squeeze(labels), t.int32)
    y, idx = t.unique(tmp_labels)
    which_rank = t.shape(y)[0] - 1
    which_rank = t.cast(which_rank, t.float32)  
    raw_pairwise_label_scores = scores.get_pairwise_scores(labels)
    mask = masks.list_mask(raw_pairwise_label_scores)
    tmp_denominator = t.matmul(mask, t.exp(predicted_scores))
    multi_label_loss = t.log(1 + tmp_denominator / t.exp(predicted_scores))
    multi_label_loss *= 2**labels - 1
    permutation_loss = t.reduce_sum(multi_label_loss)
    permutation_loss /= which_rank
    return permutation_loss
//  [python]
def get_permutation_loss(labels, predicted_scores):
    tmp_labels = tf.cast(tf.squeeze(labels), tf.int32)
    y, idx = tf.unique(tmp_labels)
    which_rank = tf.shape(y)[0] - 1
    which_rank = tf.cast(which_rank, tf.float32)  
    raw_pairwise_label_scores = scores.get_pairwise_scores(labels)
    mask = masks.list_mask(raw_pairwise_label_scores)
    tmp_denominator = tf.matmul(mask, tf.exp(predicted_scores))
    multi_label_loss = tf.log(1 + tmp_denominator / tf.exp(predicted_scores))
    multi_label_loss *= 2**labels - 1
    permutation_loss = tf.reduce_sum(multi_label_loss)
    permutation_loss /= which_rank
    return permutation_loss
//  [python]
def supervised_nt_xent_loss(z, y, temperature=0.5, base_temperature=0.07):
    '''
    Supervised normalized temperature-scaled cross entropy loss. 
    A variant of Multi-class N-pair Loss from (Sohn 2016)
    Later used in SimCLR (Chen et al. 2020, Khosla et al. 2020).
    Implementation modified from: 
        - https://github.com/google-research/simclr/blob/master/objective.py
        - https://github.com/HobbitLong/SupContrast/blob/master/losses.py
    Args:
        z: hidden vector of shape [bsz, n_features].
        y: ground truth of shape [bsz].
    '''
    batch_size = tf.shape(z)[0]
    contrast_count = 1
    anchor_count = contrast_count
    y = tf.expand_dims(y, -1)

    # mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j
    #     has the same class as sample i. Can be asymmetric.
    mask = tf.cast(tf.equal(y, tf.transpose(y)), tf.float32)
    anchor_dot_contrast = tf.divide(
        tf.matmul(z, tf.transpose(z)),
        temperature
    )
    # # for numerical stability
    logits_max = tf.reduce_max(anchor_dot_contrast, axis=1, keepdims=True)
    logits = anchor_dot_contrast - logits_max
    # # tile mask
    logits_mask = tf.ones_like(mask) - tf.eye(batch_size)
    mask = mask * logits_mask
    # compute log_prob
    exp_logits = tf.exp(logits) * logits_mask
    log_prob = logits - \
        tf.math.log(tf.reduce_sum(exp_logits, axis=1, keepdims=True))

    # compute mean of log-likelihood over positive
    # this may introduce NaNs due to zero division,
    # when a class only has one example in the batch
    mask_sum = tf.reduce_sum(mask, axis=1)
    mean_log_prob_pos = tf.reduce_sum(
        mask * log_prob, axis=1)[mask_sum > 0] / mask_sum[mask_sum > 0]

    # loss
    loss = -(temperature / base_temperature) * mean_log_prob_pos
    # loss = tf.reduce_mean(tf.reshape(loss, [anchor_count, batch_size]))
    loss = tf.reduce_mean(loss)
    return loss
//  [python]
def get_residual(labels, predicted_scores):
    tmp_labels = t.cast(t.squeeze(labels), t.int32)
    y, idx = t.unique(tmp_labels)
    which_rank = t.shape(y)[0] - 1
    which_rank = t.cast(which_rank, t.float32)
    exp_predicted_scores = t.exp(predicted_scores)
    raw_pairwise_label_scores = scores.get_pairwise_scores(labels)
    # gradient with respect to d
    d_mask = masks.list_mask(raw_pairwise_label_scores)
    tmp_denominator = t.matmul(d_mask, exp_predicted_scores)
    # loss_score_gradients
    residuals = 1 / (1 + exp_predicted_scores / tmp_denominator)
    gains = 2**labels - 1
    residuals *= gains
    # gradient with respect to d'
    _tmp_denominator = tmp_denominator + exp_predicted_scores
    _d_mask= masks.list_negative_mask(raw_pairwise_label_scores)
    _residuals = t.matmul(_d_mask, gains / _tmp_denominator)
    _residuals *= exp_predicted_scores
    residuals += _residuals
    residuals /= which_rank
    return residuals
//  [python]
def get_residual(labels, predicted_scores):
    tmp_labels = tf.cast(tf.squeeze(labels), tf.int32)
    y, idx = tf.unique(tmp_labels)
    which_rank = tf.shape(y)[0] - 1
    which_rank = tf.cast(which_rank, tf.float32)
    exp_predicted_scores = tf.exp(predicted_scores)
    raw_pairwise_label_scores = scores.get_pairwise_scores(labels)
    # gradient with respect to d
    d_mask = masks.list_mask(raw_pairwise_label_scores)
    tmp_denominator = tf.matmul(d_mask, exp_predicted_scores)
    # loss_score_gradients
    residuals = 1 / (1 + exp_predicted_scores / tmp_denominator)
    gains = 2**labels - 1
    residuals *= gains
    # gradient with respect to d'
    _tmp_denominator = tmp_denominator + exp_predicted_scores
    _d_mask= masks.list_negative_mask(raw_pairwise_label_scores)
    _residuals = tf.matmul(_d_mask, gains / _tmp_denominator)
    _residuals *= exp_predicted_scores
    residuals += _residuals
    residuals /= which_rank
    return residuals
//  [python]
def loss_layer(conv, pred, label, bboxes, stride, num_class, iou_loss_thresh):
    conv_shape = tf.shape(conv)
    batch_size = conv_shape[0]
    output_size = conv_shape[1]
    input_size = stride * output_size
    conv = tf.reshape(conv, (batch_size, output_size, output_size,
                             3, 5 + num_class))
    conv_raw_prob = conv[:, :, :, :, 5:]

    pred_xywh = pred[:, :, :, :, 0:4]
    pred_conf = pred[:, :, :, :, 4:5]

    label_xywh = label[:, :, :, :, 0:4]
    respond_bbox = label[:, :, :, :, 4:5]
    label_prob = label[:, :, :, :, 5:]

    ciou = tf.expand_dims(bbox_ciou(pred_xywh, label_xywh), axis=-1)  # (8, 13, 13, 3, 1)
    input_size = tf.cast(input_size, tf.float32)

    # 每个预测框xxxiou_loss的权重 = 2 - (ground truth的面积/图片面积)
    bbox_loss_scale = 2.0 - 1.0 * label_xywh[:, :, :, :, 2:3] * label_xywh[:, :, :, :, 3:4] / (input_size ** 2)
    ciou_loss = respond_bbox * bbox_loss_scale * (1 - ciou)  # 1. respond_bbox作为mask，有物体才计算xxxiou_loss

    # 2. respond_bbox作为mask，有物体才计算类别loss
    prob_loss = respond_bbox * tf.nn.sigmoid_cross_entropy_with_logits(labels=label_prob, logits=conv_raw_prob)

    # 3. xxxiou_loss和类别loss比较简单。重要的是conf_loss，是一个focal_loss
    # 分两步：第一步是确定 grid_h * grid_w * 3 个预测框 哪些作为反例；第二步是计算focal_loss。
    expand_pred_xywh = pred_xywh[:, :, :, :, np.newaxis, :]  # 扩展为(?, grid_h, grid_w, 3,   1, 4)
    expand_bboxes = bboxes[:, np.newaxis, np.newaxis, np.newaxis, :, :]  # 扩展为(?,      1,      1, 1, 150, 4)
    iou = bbox_iou(expand_pred_xywh, expand_bboxes)  # 所有格子的3个预测框 分别 和  150个ground truth  计算iou。   (?, grid_h, grid_w, 3, 150)
    max_iou = tf.expand_dims(tf.reduce_max(iou, axis=-1), axis=-1)  # 与150个ground truth的iou中，保留最大那个iou。  (?, grid_h, grid_w, 3, 1)

    # respond_bgd代表  这个分支输出的 grid_h * grid_w * 3 个预测框是否是 反例（背景）
    # label有物体，respond_bgd是0。 没物体的话：如果和某个gt(共150个)的iou超过iou_loss_thresh，respond_bgd是0；如果和所有gt(最多150个)的iou都小于iou_loss_thresh，respond_bgd是1。
    # respond_bgd是0代表有物体，不是反例；  权重respond_bgd是1代表没有物体，是反例。
    # 有趣的是，模型训练时由于不断更新，对于同一张图片，两次预测的 grid_h * grid_w * 3 个预测框（对于这个分支输出）  是不同的。用的是这些预测框来与gt计算iou来确定哪些预测框是反例。
    # 而不是用固定大小（不固定位置）的先验框。
    respond_bgd = (1.0 - respond_bbox) * tf.cast(max_iou < iou_loss_thresh, tf.float32)

    # 二值交叉熵损失
    pos_loss = respond_bbox * (0 - K.log(pred_conf + K.epsilon()))
    neg_loss = respond_bgd  * (0 - K.log(1 - pred_conf + K.epsilon()))

    conf_loss = pos_loss + neg_loss
    # 回顾respond_bgd，某个预测框和某个gt的iou超过iou_loss_thresh，不被当作是反例。在参与“预测的置信位 和 真实置信位 的 二值交叉熵”时，这个框也可能不是正例(label里没标这个框是1的话)。这个框有可能不参与置信度loss的计算。
    # 这种框一般是gt框附近的框，或者是gt框所在格子的另外两个框。它既不是正例也不是反例不参与置信度loss的计算。（论文里称之为ignore）

    ciou_loss = tf.reduce_mean(tf.reduce_sum(ciou_loss, axis=[1, 2, 3, 4]))  # 每个样本单独计算自己的ciou_loss，再求平均值
    conf_loss = tf.reduce_mean(tf.reduce_sum(conf_loss, axis=[1, 2, 3, 4]))  # 每个样本单独计算自己的conf_loss，再求平均值
    prob_loss = tf.reduce_mean(tf.reduce_sum(prob_loss, axis=[1, 2, 3, 4]))  # 每个样本单独计算自己的prob_loss，再求平均值

    return ciou_loss, conf_loss, prob_loss
//  [python]
def get_lambda_permutation_loss(labels, predicted_scores):
    tmp_labels = t.cast(t.squeeze(labels), t.int32)
    y, idx = t.unique(tmp_labels)
    which_rank = t.shape(y)[0] - 1
    which_rank = t.cast(which_rank, t.float32)
    raw_pairwise_label_scores = scores.get_pairwise_scores(labels)
    mask = masks.list_mask(raw_pairwise_label_scores)
    # calculate delta_z
    gains = 2**labels - 1
    raw_pairwise_scores = scores.get_pairwise_scores(predicted_scores)
    score_mask = masks.list_mask(raw_pairwise_scores)
    ranks = t.reduce_sum(score_mask, axis=1)
    ranks = t.reshape(ranks, [-1, 1])
    log_2 = t.log(t.constant(2.0, dtype=t.float32))
    cg_discounts = t.log(ranks + 2.0) / log_2
    cg_discounts = t.reshape(cg_discounts, [-1, 1])
    raw_pairwise_discounts = t.abs(scores.get_pairwise_scores(cg_discounts))
    raw_pairwise_gains = t.abs(scores.get_pairwise_scores(gains))
    delta_z = raw_pairwise_gains
    delta_z = t.multiply(raw_pairwise_discounts, raw_pairwise_gains)
    abs_pairwise_scores = t.abs(raw_pairwise_scores)
    delta_z = t.divide(delta_z, abs_pairwise_scores+1e-7)
    # max_dcg = t.reduce_sum(math_fns.cal_idcg_ks(labels, 10))
    max_dcg = t.reduce_sum(math_fns.cal_idcg_ks(labels, array_ops.size(predicted_scores)))
    delta_z = t.divide(delta_z, max_dcg)
    delta_mask = t.multiply(mask, delta_z)
    delta_mask = t.stop_gradient(delta_mask)

    tmp_denominator = t.matmul(delta_mask, t.exp(predicted_scores))
    multi_label_loss = t.log(1 + tmp_denominator / t.exp(predicted_scores))
    permutation_loss = t.reduce_sum(multi_label_loss)
    permutation_loss /= which_rank
    # # equal pairwise loss
    # pairwise_prediction_scores = scores.get_pairwise_scores(predicted_scores)
    # pairwise_equal_loss = loss_fns.get_equal_pair_loss(raw_pairwise_label_scores, \
    #     pairwise_prediction_scores)
    # permutation_loss += pairwise_equal_loss
    return permutation_loss
//  [python]
def get_lambda_permutation_loss(labels, predicted_scores):
    tmp_labels = tf.cast(tf.squeeze(labels), tf.int32)
    y, idx = tf.unique(tmp_labels)
    which_rank = tf.shape(y)[0] - 1
    which_rank = tf.cast(which_rank, tf.float32)
    raw_pairwise_label_scores = scores.get_pairwise_scores(labels)
    mask = masks.list_mask(raw_pairwise_label_scores)
    # calculate delta_z
    gains = 2**labels - 1
    raw_pairwise_scores = scores.get_pairwise_scores(predicted_scores)
    score_mask = masks.list_mask(raw_pairwise_scores)
    ranks = tf.reduce_sum(score_mask, axis=1)
    ranks = tf.reshape(ranks, [-1, 1])
    log_2 = tf.log(tf.constant(2.0, dtype=tf.float32))
    cg_discounts = tf.log(ranks + 2.0) / log_2
    cg_discounts = tf.reshape(cg_discounts, [-1, 1])
    raw_pairwise_discounts = tf.abs(scores.get_pairwise_scores(cg_discounts))
    raw_pairwise_gains = tf.abs(scores.get_pairwise_scores(gains))
    delta_z = raw_pairwise_gains
    delta_z = tf.multiply(raw_pairwise_discounts, raw_pairwise_gains)
    abs_pairwise_scores = tf.abs(raw_pairwise_scores)
    delta_z = tf.divide(delta_z, abs_pairwise_scores+1e-7)
    # max_dcg = tf.reduce_sum(math_fns.cal_idcg_ks(labels, 10))
    max_dcg = tf.reduce_sum(math_fns.cal_idcg_ks(labels, array_ops.size(predicted_scores)))
    delta_z = tf.divide(delta_z, max_dcg)
    delta_mask = tf.multiply(mask, delta_z)
    delta_mask = tf.stop_gradient(delta_mask)

    tmp_denominator = tf.matmul(delta_mask, tf.exp(predicted_scores))
    multi_label_loss = tf.log(1 + tmp_denominator / tf.exp(predicted_scores))
    permutation_loss = tf.reduce_sum(multi_label_loss)
    permutation_loss /= which_rank
    # # equal pairwise loss
    # pairwise_prediction_scores = scores.get_pairwise_scores(predicted_scores)
    # pairwise_equal_loss = loss_fns.get_equal_pair_loss(raw_pairwise_label_scores, \
    #     pairwise_prediction_scores)
    # permutation_loss += pairwise_equal_loss
    return permutation_loss
//  [python]
def _get_ur_logits(features, params):
    # mlp for rnn
    with tf.variable_scope('ur', reuse=tf.AUTO_REUSE):
        # mlp for rnn
        out_l0 = tf.layers.dense(features, params.mlp_sizes[0], name='full_dense_l0', activation=tf.nn.relu)
        out_l1 = tf.layers.dense(out_l0, params.mlp_sizes[-1], name='full_dense_l1', activation=tf.nn.relu)
        # dropout or batch normalization did not help
        # out_l1 = tf.nn.dropout(out_l1, params.dropout_rate)
        x = tf.squeeze(out_l1)
        x = tf.reshape(x, [-1, params.mlp_sizes[-1]])
        h_1 = tf.zeros(tf.shape(x))
        h_1 = tf.to_float(h_1)
        # ls, current = gru(x, h_1, params)
        if params.rnn == 'C2':
            ls, current = rnnC2(x, h_1, params)
        else:
            ls, current = rnnC1(x, h_1, params)
    return ls
//  [python]
def _get_ur_logits(features, params):
    # mlp for rnn
    with tf.variable_scope('ur', reuse=tf.AUTO_REUSE):
        # mlp for rnn
        out_l0 = tf.layers.dense(features, params.mlp_sizes[0], name='full_dense_l0', activation=tf.nn.relu)
        out_l1 = tf.layers.dense(out_l0, params.mlp_sizes[-1], name='full_dense_l1', activation=tf.nn.relu)
        # dropout or batch normalization did not help
        # out_l1 = tf.nn.dropout(out_l1, params.dropout_rate)
        x = tf.squeeze(out_l1)
        x = tf.reshape(x, [-1, params.mlp_sizes[-1]])
        h_1 = tf.zeros(tf.shape(x))
        h_1 = tf.cast(h_1, dtype=tf.float32)
        # ls, current = gru(x, h_1, params)
        if params.rnn == 'C2':
            ls, current = rnnC2(x, h_1, params)
        else:
            ls, current = rnnC1(x, h_1, params)
    return ls
//  [python]
def _get_actions_ratings(keep_labels, total_predicted_scores, iteration, is_training):
    # take all indices that give the current max label scores during training
    true_label = tf.squeeze(tf.reduce_max(keep_labels))
    ok_label_indices = tf.where(tf.equal(tf.squeeze(keep_labels), true_label))
    # take the only one index that gives the current prediction score during validation and inference
    prediction_index = tf.argmax(total_predicted_scores, axis=0)    
    if is_training:
        # keep_labels [-1, 1]
        # take all indices
        action_index = ok_label_indices
    else:
        # take the only one index
        action_index = prediction_index
    action_index = tf.to_int32(action_index)
    action_index = tf.reshape(action_index, [-1, 1])
    if is_training:
        # take all indices
        rating = tf.gather_nd(keep_labels, action_index)
        # action = rating
    else:
        # take the only one index
        # action = tf.cast(iteration, dtype=tf.float32)
        rating = tf.gather(keep_labels, action_index)
    return rating
//  [python]
def _get_actions_ratings(keep_labels, total_predicted_scores, iteration, is_training):
    # take all indices that give the current max label scores during training
    true_label = tf.squeeze(tf.reduce_max(keep_labels))
    ok_label_indices = tf.where(tf.equal(tf.squeeze(keep_labels), true_label))
    # take the only one index that gives the current prediction score during validation and inference
    prediction_index = tf.argmax(total_predicted_scores, axis=0)    
    if is_training:
        # keep_labels [-1, 1]
        # take all indices
        action_index = ok_label_indices
    else:
        # take the only one index
        action_index = prediction_index
    action_index = tf.cast(action_index, tf.int32)
    action_index = tf.reshape(action_index, [-1, 1])
    if is_training:
        # take all indices
        rating = tf.gather_nd(keep_labels, action_index)
        # action = rating
    else:
        # take the only one index
        # action = tf.cast(iteration, dtype=tf.float32)
        rating = tf.gather(keep_labels, action_index)
    return rating
//  [python]
def _get_rnn_leave_one_predictions_labels(keep_labels, keep_predicted_scores, 
    x, current, n_data, action_index, rnn_state_size):
    action_index = t.cast(action_index, t.int32)
    indices = t.range(n_data, dtype=t.int32)
    keep_indices = t.where(t.not_equal(indices, t.squeeze(action_index)))
    keep_predicted_scores = t.gather_nd(keep_predicted_scores, keep_indices)
    keep_labels = t.gather_nd(keep_labels, keep_indices)
    # keep/check the shape [-1, 1]
    keep_predicted_scores = t.reshape(keep_predicted_scores, [-1, 1])
    keep_labels = t.reshape(keep_labels, [-1, 1])
    # new x ht
    x = t.gather_nd(x, keep_indices)
    ht_1 = t.gather(current, t.reshape(action_index, [-1, 1]))
    ht_1 = t.reshape(ht_1, [1, rnn_state_size])   
    return keep_labels, keep_predicted_scores, x, ht_1
//  [python]
def _get_rnn_leave_one_predictions_labels(keep_labels, keep_predicted_scores, 
    x, current, n_data, action_index, rnn_state_size):
    action_index = tf.cast(action_index, tf.int32)
    indices = tf.range(n_data, dtype=tf.int32)
    keep_indices = tf.where(tf.not_equal(indices, tf.squeeze(action_index)))
    keep_predicted_scores = tf.gather_nd(keep_predicted_scores, keep_indices)
    keep_labels = tf.gather_nd(keep_labels, keep_indices)
    # keep/check the shape [-1, 1]
    keep_predicted_scores = tf.reshape(keep_predicted_scores, [-1, 1])
    keep_labels = tf.reshape(keep_labels, [-1, 1])
    # new x ht
    x = tf.gather_nd(x, keep_indices)
    ht_1 = tf.gather(current, tf.reshape(action_index, [-1, 1]))
    ht_1 = tf.reshape(ht_1, [1, rnn_state_size])   
    return keep_labels, keep_predicted_scores, x, ht_1
//  [python]
def get_lambda_residual(labels, predicted_scores):
    tmp_labels = t.cast(t.squeeze(labels), t.int32)
    y, idx = t.unique(tmp_labels)
    which_rank = t.shape(y)[0] - 1
    which_rank = t.cast(which_rank, t.float32)
    exp_predicted_scores = t.exp(predicted_scores)
    raw_pairwise_label_scores = scores.get_pairwise_scores(labels)
    # gradient with respect to d
    d_mask = masks.list_mask(raw_pairwise_label_scores)
    # calculate delta_z
    gains = 2**labels - 1
    raw_pairwise_scores = scores.get_pairwise_scores(predicted_scores)
    score_mask = masks.list_mask(raw_pairwise_scores)
    ranks = t.reduce_sum(score_mask, axis=1)
    ranks = t.reshape(ranks, [-1, 1])
    log_2 = t.log(t.constant(2.0, dtype=t.float32))
    cg_discounts = t.log(ranks + 2.0) / log_2
    cg_discounts = t.reshape(cg_discounts, [-1, 1])
    raw_pairwise_discounts = t.abs(scores.get_pairwise_scores(cg_discounts))
    raw_pairwise_gains = t.abs(scores.get_pairwise_scores(gains))
    delta_z = raw_pairwise_gains
    delta_z = t.multiply(raw_pairwise_discounts, raw_pairwise_gains)
    abs_pairwise_scores = t.abs(raw_pairwise_scores)
    delta_z = t.divide(delta_z, abs_pairwise_scores+1e-7)
    # max_dcg = t.reduce_sum(math_fns.cal_idcg_ks(labels, 10))
    max_dcg = t.reduce_sum(math_fns.cal_idcg_ks(labels, array_ops.size(predicted_scores)))
    delta_z = t.divide(delta_z, max_dcg)
    delta_d_mask = t.multiply(d_mask, delta_z)
    delta_d_mask = t.stop_gradient(delta_d_mask)
    tmp_denominator = t.matmul(d_mask, exp_predicted_scores)
    # loss_score_gradients
    residuals = 1 / (1 + exp_predicted_scores / tmp_denominator)
    gains = 2**labels - 1
    # gain_sum = t.reduce_sum(gains)
    # residuals /= gain_sum
    # residuals /= which_rank
    # residuals /= num_not_min
    # gradient with respect to d'
    _tmp_denominator = tmp_denominator + exp_predicted_scores
    _d_mask= masks.list_negative_mask(raw_pairwise_label_scores)

    delta__d_mask = t.multiply(_d_mask, delta_z)
    delta__d_mask = t.stop_gradient(delta__d_mask)

    _residuals = t.matmul(delta__d_mask, 1 / _tmp_denominator)
    _residuals *= exp_predicted_scores
    residuals += _residuals
    residuals /= which_rank
    return residuals
//  [python]
def get_lambda_residual(labels, predicted_scores):
    tmp_labels = tf.cast(tf.squeeze(labels), tf.int32)
    y, idx = tf.unique(tmp_labels)
    which_rank = tf.shape(y)[0] - 1
    which_rank = tf.cast(which_rank, tf.float32)
    exp_predicted_scores = tf.exp(predicted_scores)
    raw_pairwise_label_scores = scores.get_pairwise_scores(labels)
    # gradient with respect to d
    d_mask = masks.list_mask(raw_pairwise_label_scores)
    # calculate delta_z
    gains = 2**labels - 1
    raw_pairwise_scores = scores.get_pairwise_scores(predicted_scores)
    score_mask = masks.list_mask(raw_pairwise_scores)
    ranks = tf.reduce_sum(score_mask, axis=1)
    ranks = tf.reshape(ranks, [-1, 1])
    log_2 = tf.log(tf.constant(2.0, dtype=tf.float32))
    cg_discounts = tf.log(ranks + 2.0) / log_2
    cg_discounts = tf.reshape(cg_discounts, [-1, 1])
    raw_pairwise_discounts = tf.abs(scores.get_pairwise_scores(cg_discounts))
    raw_pairwise_gains = tf.abs(scores.get_pairwise_scores(gains))
    delta_z = raw_pairwise_gains
    delta_z = tf.multiply(raw_pairwise_discounts, raw_pairwise_gains)
    abs_pairwise_scores = tf.abs(raw_pairwise_scores)
    delta_z = tf.divide(delta_z, abs_pairwise_scores+1e-7)
    # max_dcg = tf.reduce_sum(math_fns.cal_idcg_ks(labels, 10))
    max_dcg = tf.reduce_sum(math_fns.cal_idcg_ks(labels, array_ops.size(predicted_scores)))
    delta_z = tf.divide(delta_z, max_dcg)
    delta_d_mask = tf.multiply(d_mask, delta_z)
    delta_d_mask = tf.stop_gradient(delta_d_mask)
    tmp_denominator = tf.matmul(d_mask, exp_predicted_scores)
    # loss_score_gradients
    residuals = 1 / (1 + exp_predicted_scores / tmp_denominator)
    gains = 2**labels - 1
    # gain_sum = tf.reduce_sum(gains)
    # residuals /= gain_sum
    # residuals /= which_rank
    # residuals /= num_not_min
    # gradient with respect to d'
    _tmp_denominator = tmp_denominator + exp_predicted_scores
    _d_mask= masks.list_negative_mask(raw_pairwise_label_scores)

    delta__d_mask = tf.multiply(_d_mask, delta_z)
    delta__d_mask = tf.stop_gradient(delta__d_mask)

    _residuals = tf.matmul(delta__d_mask, 1 / _tmp_denominator)
    _residuals *= exp_predicted_scores
    residuals += _residuals
    residuals /= which_rank
    return residuals
//  [python]
def _get_leave_one_predictions_labels(keep_labels, keep_predicted_scores, 
    x, current, n_data, action_index, rnn_state_size):
    action_index = t.cast(action_index, t.int32)
    indices = t.range(n_data, dtype=t.int32)
    keep_indices = t.where(t.not_equal(indices, t.squeeze(action_index)))
    keep_predicted_scores = t.gather_nd(keep_predicted_scores, keep_indices)
    keep_labels = t.gather_nd(keep_labels, keep_indices)
    # keep/check the shape [-1, 1]
    keep_predicted_scores = t.reshape(keep_predicted_scores, [-1, 1])
    keep_labels = t.reshape(keep_labels, [-1, 1])
    # new x st ct
    st, ct = t.unstack(current)
    x = t.gather_nd(x, keep_indices)
    st_1 = t.gather(st, t.reshape(action_index, [-1, 1]))
    st_1 = t.reshape(st_1, [1, rnn_state_size])
    ct_1 = t.gather(ct, t.reshape(action_index, [-1, 1]))
    ct_1 = t.reshape(ct_1, [1, rnn_state_size])
    prev = t.stack([st_1, ct_1])    
    return keep_labels, keep_predicted_scores, x, prev
//  [python]
def _get_leave_one_predictions_labels(keep_labels, keep_predicted_scores, 
    x, current, n_data, action_index, rnn_state_size):
    action_index = tf.cast(action_index, tf.int32)
    indices = tf.range(n_data, dtype=tf.int32)
    keep_indices = tf.where(tf.not_equal(indices, tf.squeeze(action_index)))
    keep_predicted_scores = tf.gather_nd(keep_predicted_scores, keep_indices)
    keep_labels = tf.gather_nd(keep_labels, keep_indices)
    # keep/check the shape [-1, 1]
    keep_predicted_scores = tf.reshape(keep_predicted_scores, [-1, 1])
    keep_labels = tf.reshape(keep_labels, [-1, 1])
    # new x st ct
    st, ct = tf.unstack(current)
    x = tf.gather_nd(x, keep_indices)
    st_1 = tf.gather(st, tf.reshape(action_index, [-1, 1]))
    st_1 = tf.reshape(st_1, [1, rnn_state_size])
    ct_1 = tf.gather(ct, tf.reshape(action_index, [-1, 1]))
    ct_1 = tf.reshape(ct_1, [1, rnn_state_size])
    prev = tf.stack([st_1, ct_1])    
    return keep_labels, keep_predicted_scores, x, prev
//  [python]
def convert_binary_array_to_integer(binary_tensor):
  """Convert batched binary tensor to integers."""
  # The input binary tensor must be left-padded to the same length.
  # e.g. [[0, 0, 0, 0, 1, 1, 1],
  #       [0, 1, 0, 0, 1, 0, 0],
  #       [1, 1, 0, 0, 1, 0, 0]]
  # returns [7, 36, 100]
  binary_tensor = tf.convert_to_tensor(binary_tensor)
  return tf.reduce_sum(
      tf.reverse(tensor=binary_tensor, axis=[-1]) * 2 ** tf.range(tf.shape(binary_tensor)[-1]),
      axis=-1)
//  [python]
def convert_binary_array_to_integer(binary_tensor):
  """Convert batched binary tensor to integers."""
  # The input binary tensor must be left-padded to the same length.
  # e.g. [[0, 0, 0, 0, 1, 1, 1],
  #       [0, 1, 0, 0, 1, 0, 0],
  #       [1, 1, 0, 0, 1, 0, 0]]
  # returns [7, 36, 100]
  binary_tensor = tf.convert_to_tensor(binary_tensor)
  return tf.reduce_sum(
      tf.cast(tf.reverse(tensor=binary_tensor, axis=[-1]), dtype=tf.int32)
      * 2 ** tf.range(tf.cast(tf.shape(binary_tensor)[-1], dtype=tf.int32)),
      axis=-1)
//  [python]
def __call__(self,
               input_tensor,
               samples=None,
               num_samples=()):
    dist = self.networks(input_tensor)
    if samples is None:
      samples = dist.sample(num_samples)
    samples = tf.to_float(samples)
    likelihood = dist.log_prob(samples)
    return samples, likelihood
//  [python]
def __call__(self,
               input_tensor,
               samples=None,
               num_samples=()):
    dist = self.networks(input_tensor)
    if samples is None:
      samples = dist.sample(num_samples)
    samples = tf.cast(samples, tf.float32)
    likelihood = dist.log_prob(samples)
    return samples, likelihood
//  [python]
def __call__(self,
               input_tensor,
               samples=None,
               num_samples=(),
               stick_breaking=False):
    if stick_breaking:
      raise NotImplementedError(
          'When using stick breaking, please use the get_logits, '
          'get_stick_breaking_samples function to generate the samples, '
          'and calculate the elbo.')

    else:
      logits = self.get_logits(input_tensor)
      dist = tfd.Categorical(logits=logits)
      if samples is None:
        samples = dist.sample(num_samples)
      samples = t.cast(samples, tf.float32)
      likelihood = dist.log_prob(samples)
      return samples, likelihood, logits
//  [python]
def __call__(self,
               input_tensor,
               samples=None,
               num_samples=(),
               stick_breaking=False):
    if stick_breaking:
      raise NotImplementedError(
          'When using stick breaking, please use the get_logits, '
          'get_stick_breaking_samples function to generate the samples, '
          'and calculate the elbo.')

    else:
      logits = self.get_logits(input_tensor)
      dist = tfd.Categorical(logits=logits)
      if samples is None:
        samples = dist.sample(num_samples)
      samples = tf.cast(samples, tf.float32)
      likelihood = dist.log_prob(samples)
      return samples, likelihood, logits
//  [python]
def get_stick_breaking_samples(encoder_logits, num_samples=1,
                               u_noise=None, tree_structure=False):
  """Generate categorical samples with stick breaking construction."""
  # in the case of stick breaking, the encoder output K-1 category logits
  batch_size, num_variables, num_categories = encoder_logits.shape

  sigma_phi = tf.math.sigmoid(encoder_logits)
  if u_noise is None:
    # [num_samples * batch_size, num_variables, num_categories]
    u_noise = sample_uniform_random_variable(
        [batch_size, num_variables, num_categories],
        num_samples=num_samples)
  sigma_phi_tiled = tf.tile(sigma_phi, [num_samples, 1, 1])
  b = tf.to_float(tf.less(u_noise, sigma_phi_tiled))

  b = tf.reshape(b, [-1, num_categories])
  # set the label of the K-th category to be 1.
  b_extended = tf.concat(
      [b, tf.ones_like(b[:, 0], dtype=tf.float32)[:, tf.newaxis]],
      axis=-1)

  if tree_structure:
    # logits_mask masks the nodes along the path of the tree
    categorical_samples, logits_mask = convert_balanced_tree_to_category(b)
  else:
    # logits_mask the first appearance of 1
    categorical_samples, logits_mask = get_first_appearance_in_2d_tensor(
        b_extended, target=1, with_mask=True)
  categorical_samples = tf.reshape(
      categorical_samples, [num_samples, batch_size, num_variables])
  logits_mask = tf.reshape(
      logits_mask,
      [num_samples, batch_size, num_variables, num_categories+1])

  binary_samples = tf.reshape(
      b_extended, [num_samples, batch_size, num_variables, -1])

  return categorical_samples, logits_mask, u_noise, binary_samples
//  [python]
def get_stick_breaking_samples(encoder_logits, num_samples=1,
                               u_noise=None, tree_structure=False):
  """Generate categorical samples with stick breaking construction."""
  # in the case of stick breaking, the encoder output K-1 category logits
  batch_size, num_variables, num_categories = encoder_logits.shape

  sigma_phi = tf.math.sigmoid(encoder_logits)
  if u_noise is None:
    # [num_samples * batch_size, num_variables, num_categories]
    u_noise = sample_uniform_random_variable(
        [batch_size, num_variables, num_categories],
        num_samples=num_samples)
  sigma_phi_tiled = tf.tile(sigma_phi, [num_samples, 1, 1])
  b = tf.cast(u_noise < sigma_phi_tiled, tf.float32)

  b = tf.reshape(b, [-1, num_categories])
  # set the label of the K-th category to be 1.
  b_extended = tf.concat(
      [b, tf.ones_like(b[:, 0], dtype=tf.float32)[:, tf.newaxis]],
      axis=-1)

  if tree_structure:
    # logits_mask masks the nodes along the path of the tree
    categorical_samples, logits_mask = convert_balanced_tree_to_category(b)
  else:
    # logits_mask the first appearance of 1
    categorical_samples, logits_mask = get_first_appearance_in_2d_tensor(
        b_extended, target=1, with_mask=True)
  categorical_samples = tf.reshape(
      categorical_samples, [num_samples, batch_size, num_variables])
  logits_mask = tf.reshape(
      logits_mask,
      [num_samples, batch_size, num_variables, num_categories+1])

  binary_samples = tf.reshape(
      b_extended, [num_samples, batch_size, num_variables, -1])

  return categorical_samples, logits_mask, u_noise, binary_samples
//  [python]
def __call__(self,
               input_tensor,
               samples=None,
               num_samples=()):
    logits = self.get_logits(input_tensor)
    dist = tfd.Bernoulli(logits=logits)
    if samples is None:
      samples = dist.sample(num_samples)
    samples = t.cast(samples, tf.float32)
    likelihood = dist.log_prob(samples)
    return samples, likelihood, logits
//  [python]
def __call__(self,
               input_tensor,
               samples=None,
               num_samples=()):
    logits = self.get_logits(input_tensor)
    dist = tfd.Bernoulli(logits=logits)
    if samples is None:
      samples = dist.sample(num_samples)
    samples = tf.cast(samples, tf.float32)
    likelihood = dist.log_prob(samples)
    return samples, likelihood, logits
//  [python]
def __call__(self,
               input_tensor,
               samples=None,
               num_samples=(),
               stick_breaking=False):
    if stick_breaking:
      # if stick_breaking we will convert the logits from K categories to K-1
      # and append a non-trainable zero logits to K-th categories
      raise NotImplementedError(
          'When using stick breaking, please use the get_logits, '
          'get_stick_breaking_samples function to generate the samples, '
          'and calculate the elbo.')

    else:
      logits = self.get_logits(input_tensor)
      dist = tfd.Categorical(logits=logits)
      if samples is None:
        samples = dist.sample(num_samples)
      samples = tf.cast(samples, tf.float32)
      likelihood = dist.log_prob(samples)
      return samples, likelihood, logits
//  [python]
def __init__(self,
               hidden_sizes,
               activations,
               mean_xs=None,
               demean_input=False,
               final_layer_bias_initializer='zeros',
               name='binarynet',
               kernel_initializer='glorot_uniform'):

    super(BinaryNetwork, self).__init__(name=name)
    assert len(activations) == len(hidden_sizes)

    num_layers = len(hidden_sizes)
    self.hidden_sizes = hidden_sizes
    self.activations = activations
    self.networks = keras.Sequential()

    self.networks.add(keras.layers.Flatten())

    if demean_input:
      if mean_xs is not None:
        self.networks.add(
            tf.keras.layers.Lambda(lambda x: x - mean_xs))
      else:
        self.networks.add(
            tf.keras.layers.Lambda(lambda x: 2.*tf.cast(x, tf.float32) - 1.))
    for i in range(num_layers-1):
      self.networks.add(
          keras.layers.Dense(
              units=hidden_sizes[i],
              activation=activations[i],
              kernel_initializer=kernel_initializer))

    self.networks.add(
        keras.layers.Dense(
            units=hidden_sizes[-1],
            activation=activations[-1],
            kernel_initializer=kernel_initializer,
            bias_initializer=final_layer_bias_initializer))
//  [python]
def __init__(self,
               hidden_sizes,
               activations,
               num_categories,
               mean_xs=None,
               demean_input=False,
               name='categnet',
               kernel_initializer='glorot_uniform'):

    super(CategoricalNetwork, self).__init__(name=name)
    assert len(activations) == len(hidden_sizes)

    num_layers = len(hidden_sizes)
    self.hidden_sizes = hidden_sizes
    self.activations = activations
    self.networks = keras.Sequential()
    self.num_categories = num_categories

    if demean_input:
      if mean_xs is not None:
        self.networks.add(
            tf.keras.layers.Lambda(lambda x: x - mean_xs))
      else:
        self.networks.add(
            tf.keras.layers.Lambda(lambda x: 2.*tf.cast(x, tf.float32) - 1.))
    for i in range(num_layers-1):
      self.networks.add(
          keras.layers.Dense(
              units=hidden_sizes[i],
              activation=activations[i],
              kernel_initializer=kernel_initializer))

    self.networks.add(
        keras.layers.Dense(
            units=hidden_sizes[-1] * self.num_categories,
            activation=activations[-1],
            kernel_initializer=kernel_initializer))
//  [python]
def categorical_call(self, input_tensor, hidden_samples=None, num_samples=1):
    batch_size = input_tensor.shape[0]

    hidden_sample, encoder_llk, encoder_logits = self.encoder(
        input_tensor,
        samples=hidden_samples,
        num_samples=num_samples)

    log_pb = self.prior_dist.log_prob(hidden_sample)

    if self.one_hot_sample:
      one_hot_sample = tf.one_hot(tf.to_int32(hidden_sample),
                                  depth=self.num_categories,
                                  dtype=tf.float32)
      one_hot_sample = tf.reshape(one_hot_sample,
                                  [num_samples * batch_size, -1])

      decoder_llk = self.decoder(one_hot_sample, input_tensor)[1]
    else:
      hidden_sample = tf.reshape(hidden_sample,
                                 [num_samples * batch_size, -1])
      decoder_llk = self.decoder(hidden_sample, input_tensor)[1]

    decoder_llk = tf.reshape(decoder_llk, [num_samples * batch_size, -1])
    elbo = (
        tf.reduce_sum(decoder_llk, axis=-1)
        + tf.reduce_sum(log_pb - encoder_llk, axis=-1))

    return elbo, hidden_sample, encoder_logits, encoder_llk
//  [python]
def categorical_call(self, input_tensor, hidden_samples=None, num_samples=1):
    batch_size = input_tensor.shape[0]

    hidden_sample, encoder_llk, encoder_logits = self.encoder(
        input_tensor,
        samples=hidden_samples,
        num_samples=num_samples)

    log_pb = self.prior_dist.log_prob(hidden_sample)

    if self.one_hot_sample:
      one_hot_sample = tf.one_hot(tf.cast(hidden_sample, tf.int32),
                                  depth=self.num_categories,
                                  dtype=tf.float32)
      one_hot_sample = tf.reshape(one_hot_sample,
                                  [num_samples * batch_size, -1])

      decoder_llk = self.decoder(one_hot_sample, input_tensor)[1]
    else:
      hidden_sample = tf.reshape(hidden_sample,
                                 [num_samples * batch_size, -1])
      decoder_llk = self.decoder(hidden_sample, input_tensor)[1]

    decoder_llk = tf.reshape(decoder_llk, [num_samples * batch_size, -1])
    elbo = (
        tf.reduce_sum(decoder_llk, axis=-1)
        + tf.reduce_sum(log_pb - encoder_llk, axis=-1))

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

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

    layer_grad = tf.matmul(tilde_dirichlet, tilde_loss)
    return layer_grad