jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def call(self, inputs, training=True):

        #(n_batch, n_points, n_features)
        point_embedding = inputs

        n_batches = tf.shape(point_embedding)[0]
        n_points = tf.shape(point_embedding)[1]

        #cannot concat sparse tensors directly as that incorrectly destroys the gradient, see
        #https://github.com/tensorflow/tensorflow/blob/df3a3375941b9e920667acfe72fb4c33a8f45503/tensorflow/python/ops/sparse_grad.py#L33
        #therefore, for training, we implement sparse concatenation by hand 
        indices_all = []
        values_all = []

        def func(args):
            ibatch, points_batch = args[0], args[1]
            dm = self.construct_sparse_dm_batch(points_batch)
            inds = tf.concat([tf.expand_dims(tf.to_int64(ibatch*tf.ones(tf.shape(dm.indices)[0], dtype=tf.int64)), -1), dm.indices], axis=-1)
            vals = dm.values
            return inds, vals

        elems = (tf.range(0, tf.to_int64(n_batches), delta=1, dtype=tf.int64), point_embedding)
        ret = tf.map_fn(func, elems, fn_output_signature=(tf.int64, tf.float32), parallel_iterations=1)
        shp = tf.shape(ret[0])
        # #now create a new SparseTensor that is a concatenation of the previous ones
        dms = tf.SparseTensor(
            tf.reshape(ret[0], (shp[0]*shp[1], shp[2])),
            tf.reshape(ret[1], (shp[0]*shp[1],)),
            (n_batches, n_points, n_points)
        )

        return tf.sparse.reorder(dms)
//  [python]
def call(self, inputs, training=True):

        #(n_batch, n_points, n_features)
        point_embedding = inputs

        n_batches = tf.shape(point_embedding)[0]
        n_points = tf.shape(point_embedding)[1]

        #cannot concat sparse tensors directly as that incorrectly destroys the gradient, see
        #https://github.com/tensorflow/tensorflow/blob/df3a3375941b9e920667acfe72fb4c33a8f45503/tensorflow/python/ops/sparse_grad.py#L33
        #therefore, for training, we implement sparse concatenation by hand 
        indices_all = []
        values_all = []

        def func(args):
            ibatch, points_batch = args[0], args[1]
            dm = self.construct_sparse_dm_batch(points_batch)
            inds = tf.concat([tf.expand_dims(tf.cast(ibatch, tf.int64)*tf.ones(tf.shape(dm.indices)[0], dtype=tf.int64), -1), dm.indices], axis=-1)
            vals = dm.values
            return inds, vals

        elems = (tf.range(0, tf.cast(n_batches, tf.int64), delta=1, dtype=tf.int64), point_embedding)
        ret = tf.map_fn(func, elems, fn_output_signature=(tf.int64, tf.float32), parallel_iterations=1)
        shp = tf.shape(ret[0])
        # #now create a new SparseTensor that is a concatenation of the previous ones
        dms = tf.SparseTensor(
            tf.reshape(ret[0], (shp[0]*shp[1], shp[2])),
            tf.reshape(ret[1], (shp[0]*shp[1],)),
            (n_batches, n_points, n_points)
        )

        return tf.sparse.reorder(dms)
//  [python]
def cls_130(y_true, y_pred):
    pred_id_onehot, pred_charge, pred_momentum = separate_prediction(y_pred)
    pred_id = t.cast(t.argmax(pred_id_onehot, axis=-1), t.int32)
    true_id, true_charge, true_momentum = separate_truth(y_true)

    msk_true = true_id[:, :, 0] == class_labels.index(130)
    msk_pos = pred_id == class_labels.index(130)
    num_true_pos = t.reduce_sum(t.cast(msk_true&msk_pos, t.float32))
    num_true = t.reduce_sum(t.cast(msk_true, t.float32))
    return num_true_pos/num_true
//  [python]
def cls_130(y_true, y_pred):
    pred_id_onehot, pred_charge, pred_momentum = separate_prediction(y_pred)
    pred_id = tf.cast(tf.argmax(pred_id_onehot, axis=-1), tf.int32)
    true_id, true_charge, true_momentum = separate_truth(y_true)

    msk_true = true_id[:, :, 0] == class_labels.index(130)
    msk_pos = pred_id == class_labels.index(130)
    num_true_pos = tf.reduce_sum(tf.cast(msk_true&msk_pos, tf.float32))
    num_true = tf.reduce_sum(tf.cast(msk_true, tf.float32))
    return num_true_pos/num_true
//  [python]
def call(self, inputs, training=True):
        X = t.cast(inputs, t.float32)
        msk_input = t.expand_dims(t.cast(X[:, :, 0] != 0, t.float32), -1)

        enc = self.enc(inputs)

        #embed inputs for graph structure prediction
        embedding_attention = self.layer_embedding(enc)
        if self.embedding_dropout:
            embedding_attention = self.embedding_dropout(embedding_attention, training)

        #create graph structure by predicting a sparse distance matrix
        dms = [dist(embedding_attention, training) for dist in self.dists]
        dm = self.addsparse(dms)

        #run graph net for multiclass id prediction
        x_id = self.gnn_id(enc, dm, training)
        to_decode = t.concat([enc, x_id], axis=-1)
        out_id_logits = self.layer_id(to_decode)
        out_charge = self.layer_charge(to_decode)

        #run graph net for regression output prediction, taking as an additonal input the ID predictions
        x_reg = self.gnn_reg(t.concat([enc, out_id_logits, out_charge], axis=-1), dm, training)
        to_decode = t.concat([enc, x_reg], axis=-1)
        pred_corr = self.layer_momentum(to_decode)

        #soft-mask elements for which the id prediction was 0  
        probabilistic_mask_good = 1.0 - t.keras.activations.softmax(out_id_logits)[:, :, 0]

        out_momentum_eta = X[:, :, 2] + pred_corr[:, :, 0]
        out_momentum_phi = X[:, :, 3] + pred_corr[:, :, 1] 
        out_momentum_E = X[:, :, 4] + pred_corr[:, :, 2]

        out_momentum = t.stack([
            out_momentum_eta * probabilistic_mask_good,
            out_momentum_phi * probabilistic_mask_good,
            out_momentum_E * probabilistic_mask_good,
        ], axis=-1)

        ret = t.concat([out_id_logits, out_momentum, out_charge], axis=-1)*msk_input
        return ret
//  [python]
def call(self, inputs, training=True):
        X = tf.cast(inputs, tf.float32)
        msk_input = tf.expand_dims(tf.cast(X[:, :, 0] != 0, tf.float32), -1)

        enc = self.enc(inputs)

        #embed inputs for graph structure prediction
        embedding_attention = self.layer_embedding(enc)
        if self.embedding_dropout:
            embedding_attention = self.embedding_dropout(embedding_attention, training)

        #create graph structure by predicting a sparse distance matrix
        dms = [dist(embedding_attention, training) for dist in self.dists]
        dm = self.addsparse(dms)

        #run graph net for multiclass id prediction
        x_id = self.gnn_id(enc, dm, training)
        to_decode = tf.concat([enc, x_id], axis=-1)
        out_id_logits = self.layer_id(to_decode)
        out_charge = self.layer_charge(to_decode)

        #run graph net for regression output prediction, taking as an additonal input the ID predictions
        x_reg = self.gnn_reg(tf.concat([enc, out_id_logits, out_charge], axis=-1), dm, training)
        to_decode = tf.concat([enc, x_reg], axis=-1)
        pred_corr = self.layer_momentum(to_decode)

        #soft-mask elements for which the id prediction was 0  
        probabilistic_mask_good = 1.0 - tf.keras.activations.softmax(out_id_logits)[:, :, 0]

        out_momentum_eta = X[:, :, 2] + pred_corr[:, :, 0]
        out_momentum_phi = X[:, :, 3] + pred_corr[:, :, 1] 
        out_momentum_E = X[:, :, 4] + pred_corr[:, :, 2]

        out_momentum = tf.stack([
            out_momentum_eta * probabilistic_mask_good,
            out_momentum_phi * probabilistic_mask_good,
            out_momentum_E * probabilistic_mask_good,
        ], axis=-1)

        ret = tf.concat([out_id_logits, out_momentum, out_charge], axis=-1)*msk_input
        return ret
//  [python]
def cls_211(y_true, y_pred):
    pred_id_onehot, pred_charge, pred_momentum = separate_prediction(y_pred)
    pred_id = t.cast(t.argmax(pred_id_onehot, axis=-1), t.int32)
    true_id, true_charge, true_momentum = separate_truth(y_true)

    msk_true = true_id[:, :, 0] == class_labels.index(211)
    msk_pos = pred_id == class_labels.index(211)
    num_true_pos = t.reduce_sum(t.cast(msk_true&msk_pos, t.float32))
    num_true = t.reduce_sum(t.cast(msk_true, t.float32))

    return num_true_pos/num_true
//  [python]
def cls_211(y_true, y_pred):
    pred_id_onehot, pred_charge, pred_momentum = separate_prediction(y_pred)
    pred_id = tf.cast(tf.argmax(pred_id_onehot, axis=-1), tf.int32)
    true_id, true_charge, true_momentum = separate_truth(y_true)

    msk_true = true_id[:, :, 0] == class_labels.index(211)
    msk_pos = pred_id == class_labels.index(211)
    num_true_pos = tf.reduce_sum(tf.cast(msk_true&msk_pos, tf.float32))
    num_true = tf.reduce_sum(tf.cast(msk_true, tf.float32))

    return num_true_pos/num_true
//  [python]
def my_loss_full(y_true, y_pred):
    pred_id_logits, pred_charge, pred_momentum = separate_prediction(y_pred)
    pred_id = t.cast(t.argmax(pred_id_logits, axis=-1), t.int32)
    true_id, true_charge, true_momentum = separate_truth(y_true)
    true_id_onehot = tf.one_hot(t.cast(true_id, t.int32), depth=len(class_labels))
    
    l1 = mult_classification_loss*tf.nn.softmax_cross_entropy_with_logits(true_id_onehot, pred_id_logits)
  
    l2_0 = mult_eta_loss*mse_unreduced(true_momentum[:, :, 0], pred_momentum[:, :, 0])
    l2_1 = mult_phi_loss*mse_unreduced(tf.math.floormod(true_momentum[:, :, 1] - pred_momentum[:, :, 1] + np.pi, 2*np.pi) - np.pi, 0.0)
    l2_2 = mult_energy_loss*mse_unreduced(true_momentum[:, :, 2], pred_momentum[:, :, 2])

    l2 = (l2_0 + l2_1 + l2_2)

    l3 = mult_charge_loss*mse_unreduced(true_charge, pred_charge)[:, :, 0]
    loss = l1 + l2 + l3

    return mult_total_loss*loss
//  [python]
def my_loss_full(y_true, y_pred):
    pred_id_logits, pred_charge, pred_momentum = separate_prediction(y_pred)
    pred_id = tf.cast(tf.argmax(pred_id_logits, axis=-1), tf.int32)
    true_id, true_charge, true_momentum = separate_truth(y_true)
    true_id_onehot = tf.one_hot(tf.cast(true_id, tf.int32), depth=len(class_labels))
    
    l1 = mult_classification_loss*tf.nn.softmax_cross_entropy_with_logits(true_id_onehot, pred_id_logits)
  
    l2_0 = mult_eta_loss*mse_unreduced(true_momentum[:, :, 0], pred_momentum[:, :, 0])
    l2_1 = mult_phi_loss*mse_unreduced(tf.math.floormod(true_momentum[:, :, 1] - pred_momentum[:, :, 1] + np.pi, 2*np.pi) - np.pi, 0.0)
    l2_2 = mult_energy_loss*mse_unreduced(true_momentum[:, :, 2], pred_momentum[:, :, 2])

    l2 = (l2_0 + l2_1 + l2_2)

    l3 = mult_charge_loss*mse_unreduced(true_charge, pred_charge)[:, :, 0]
    loss = l1 + l2 + l3

    return mult_total_loss*loss
//  [python]
def cls_22(y_true, y_pred):
    pred_id_onehot, pred_charge, pred_momentum = separate_prediction(y_pred)
    pred_id = t.cast(t.argmax(pred_id_onehot, axis=-1), t.int32)
    true_id, true_charge, true_momentum = separate_truth(y_true)

    msk_true = true_id[:, :, 0] == class_labels.index(22)
    msk_pos = pred_id == class_labels.index(22)
    num_true_pos = t.reduce_sum(t.cast(msk_true&msk_pos, t.float32))
    num_true = t.reduce_sum(t.cast(msk_true, t.float32))

    return num_true_pos/num_true
//  [python]
def cls_22(y_true, y_pred):
    pred_id_onehot, pred_charge, pred_momentum = separate_prediction(y_pred)
    pred_id = tf.cast(tf.argmax(pred_id_onehot, axis=-1), tf.int32)
    true_id, true_charge, true_momentum = separate_truth(y_true)

    msk_true = true_id[:, :, 0] == class_labels.index(22)
    msk_pos = pred_id == class_labels.index(22)
    num_true_pos = tf.reduce_sum(tf.cast(msk_true&msk_pos, tf.float32))
    num_true = tf.reduce_sum(tf.cast(msk_true, tf.float32))

    return num_true_pos/num_true
//  [python]
def cls_11(y_true, y_pred):
    pred_id_onehot, pred_charge, pred_momentum = separate_prediction(y_pred)
    pred_id = t.cast(t.argmax(pred_id_onehot, axis=-1), t.int32)
    true_id, true_charge, true_momentum = separate_truth(y_true)

    msk_true = true_id[:, :, 0] == class_labels.index(11)
    msk_pos = pred_id == class_labels.index(11)
    num_true_pos = t.reduce_sum(t.cast(msk_true&msk_pos, t.float32))
    num_true = t.reduce_sum(t.cast(msk_true, t.float32))

    return num_true_pos/num_true
//  [python]
def cls_11(y_true, y_pred):
    pred_id_onehot, pred_charge, pred_momentum = separate_prediction(y_pred)
    pred_id = tf.cast(tf.argmax(pred_id_onehot, axis=-1), tf.int32)
    true_id, true_charge, true_momentum = separate_truth(y_true)

    msk_true = true_id[:, :, 0] == class_labels.index(11)
    msk_pos = pred_id == class_labels.index(11)
    num_true_pos = tf.reduce_sum(tf.cast(msk_true&msk_pos, tf.float32))
    num_true = tf.reduce_sum(tf.cast(msk_true, tf.float32))

    return num_true_pos/num_true
//  [python]
def cls_13(y_true, y_pred):
    pred_id_onehot, pred_charge, pred_momentum = separate_prediction(y_pred)
    pred_id = t.cast(t.argmax(pred_id_onehot, axis=-1), t.int32)
    true_id, true_charge, true_momentum = separate_truth(y_true)

    msk_true = true_id[:, :, 0] == class_labels.index(13)
    msk_pos = pred_id == class_labels.index(13)
    num_true_pos = t.reduce_sum(t.cast(msk_true&msk_pos, t.float32))
    num_true = t.reduce_sum(t.cast(msk_true, t.float32))

    return num_true_pos/num_true
//  [python]
def cls_13(y_true, y_pred):
    pred_id_onehot, pred_charge, pred_momentum = separate_prediction(y_pred)
    pred_id = tf.cast(tf.argmax(pred_id_onehot, axis=-1), tf.int32)
    true_id, true_charge, true_momentum = separate_truth(y_true)

    msk_true = true_id[:, :, 0] == class_labels.index(13)
    msk_pos = pred_id == class_labels.index(13)
    num_true_pos = tf.reduce_sum(tf.cast(msk_true&msk_pos, tf.float32))
    num_true = tf.reduce_sum(tf.cast(msk_true, tf.float32))

    return num_true_pos/num_true
//  [python]
def num_pred(y_true, y_pred):
    pred_id_onehot, pred_charge, pred_momentum = separate_prediction(y_pred)
    pred_id = t.cast(t.argmax(pred_id_onehot, axis=-1), t.int32)
    true_id, true_charge, true_momentum = separate_truth(y_true)

    ntrue = t.reduce_sum(t.cast(true_id[:, :, 0]!=0, t.int32))
    npred = t.reduce_sum(t.cast(pred_id!=0, t.int32))
    return t.cast(ntrue - npred, t.float32)
//  [python]
def num_pred(y_true, y_pred):
    pred_id_onehot, pred_charge, pred_momentum = separate_prediction(y_pred)
    pred_id = tf.cast(tf.argmax(pred_id_onehot, axis=-1), tf.int32)
    true_id, true_charge, true_momentum = separate_truth(y_true)

    ntrue = tf.reduce_sum(tf.cast(true_id[:, :, 0]!=0, tf.int32))
    npred = tf.reduce_sum(tf.cast(pred_id!=0, tf.int32))
    return tf.cast(ntrue - npred, tf.float32)
//  [python]
def eta_resolution(y_true, y_pred):
    pred_id_onehot, pred_charge, pred_momentum = separate_prediction(y_pred)
    pred_id = t.cast(t.argmax(pred_id_onehot, axis=-1), t.int32)
    true_id, true_charge, true_momentum = separate_truth(y_true)

    msk = true_id[:, :, 0]!=0
    return t.reduce_mean(mse_unreduced(true_momentum[msk][:, 0], pred_momentum[msk][:, 0]))
//  [python]
def eta_resolution(y_true, y_pred):
    pred_id_onehot, pred_charge, pred_momentum = separate_prediction(y_pred)
    pred_id = tf.cast(tf.argmax(pred_id_onehot, axis=-1), tf.int32)
    true_id, true_charge, true_momentum = separate_truth(y_true)

    msk = true_id[:, :, 0]!=0
    return tf.reduce_mean(mse_unreduced(true_momentum[msk][:, 0], pred_momentum[msk][:, 0]))
//  [python]
def accuracy(y_true, y_pred):
    pred_id_onehot, pred_charge, pred_momentum = separate_prediction(y_pred)
    pred_id = t.cast(t.argmax(pred_id_onehot, axis=-1), t.int32)
    true_id, true_charge, true_momentum = separate_truth(y_true)

    is_true = true_id[:, :, 0]!=0
    is_same = true_id[:, :, 0] == pred_id

    acc = t.reduce_sum(t.cast(is_true&is_same, t.int32)) / t.reduce_sum(t.cast(is_true, t.int32))
    return t.cast(acc, t.float32)
//  [python]
def accuracy(y_true, y_pred):
    pred_id_onehot, pred_charge, pred_momentum = separate_prediction(y_pred)
    pred_id = tf.cast(tf.argmax(pred_id_onehot, axis=-1), tf.int32)
    true_id, true_charge, true_momentum = separate_truth(y_true)

    is_true = true_id[:, :, 0]!=0
    is_same = true_id[:, :, 0] == pred_id

    acc = tf.reduce_sum(tf.cast(is_true&is_same, tf.int32)) / tf.reduce_sum(tf.cast(is_true, tf.int32))
    return tf.cast(acc, tf.float32)
//  [python]
def energy_resolution(y_true, y_pred):
    pred_id_onehot, pred_charge, pred_momentum = separate_prediction(y_pred)
    pred_id = t.cast(t.argmax(pred_id_onehot, axis=-1), t.int32)
    true_id, true_charge, true_momentum = separate_truth(y_true)

    msk = true_id[:, :, 0]!=0
    return t.reduce_mean(mse_unreduced(true_momentum[msk][:, 2], pred_momentum[msk][:, 2]))
//  [python]
def energy_resolution(y_true, y_pred):
    pred_id_onehot, pred_charge, pred_momentum = separate_prediction(y_pred)
    pred_id = tf.cast(tf.argmax(pred_id_onehot, axis=-1), tf.int32)
    true_id, true_charge, true_momentum = separate_truth(y_true)

    msk = true_id[:, :, 0]!=0
    return tf.reduce_mean(mse_unreduced(true_momentum[msk][:, 2], pred_momentum[msk][:, 2]))
//  [python]
def encode(self, text, shape=None):
    """Tokenizes a text example."""
    tokens = self.tokenizer.tokenize(text)

    if shape is not None:
      tokens.set_shape(shape)

    tokens = tokens - 97
    mask = tf.cast(tokens < 0, tf.int32)
    tokens = mask * tf.constant(
        value=26, shape=tokens.shape) + (1 - mask) * tokens
    return tokens
//  [python]
def phi_resolution(y_true, y_pred):
    pred_id_onehot, pred_charge, pred_momentum = separate_prediction(y_pred)
    pred_id = t.cast(t.argmax(pred_id_onehot, axis=-1), t.int32)
    true_id, true_charge, true_momentum = separate_truth(y_true)

    msk = true_id[:, :, 0]!=0
    return t.reduce_mean(mse_unreduced(t.math.floormod(true_momentum[msk][:, 1] - pred_momentum[msk][:, 1] + np.pi, 2*np.pi) - np.pi, 0.0))
//  [python]
def phi_resolution(y_true, y_pred):
    pred_id_onehot, pred_charge, pred_momentum = separate_prediction(y_pred)
    pred_id = tf.cast(tf.argmax(pred_id_onehot, axis=-1), tf.int32)
    true_id, true_charge, true_momentum = separate_truth(y_true)

    msk = true_id[:, :, 0]!=0
    return tf.reduce_mean(mse_unreduced(tf.math.floormod(true_momentum[msk][:, 1] - pred_momentum[msk][:, 1] + np.pi, 2*np.pi) - np.pi, 0.0))
//  [python]
def decode(self, tokens):
    mask = t.cast(tokens == 26, t.int32)
    extra_mask = t.cast(tokens > 26, t.int32)

    tokens = tokens + 97
    tokens = mask * t.constant(32, shape=tokens.shape) + (1 - mask) * tokens
    tokens = extra_mask * t.constant(
        45, shape=tokens.shape) + (1 - extra_mask) * tokens

    return self.tokenizer.detokenize(tokens)
//  [python]
def decode(self, tokens):
    mask = tf.cast(tokens == 26, tf.int32)
    extra_mask = tf.cast(tokens > 26, tf.int32)

    tokens = tokens + 97
    tokens = mask * tf.constant(32, shape=tokens.shape) + (1 - mask) * tokens
    tokens = extra_mask * tf.constant(
        45, shape=tokens.shape) + (1 - extra_mask) * tokens

    return self.tokenizer.detokenize(tokens)
//  [python]
def call(self, inputs, inputs_positions=None):
    del inputs_positions
    return tf.to_float(self.gamma) * inputs
//  [python]
def call(self, inputs, inputs_positions=None):
    del inputs_positions
    return tf.cast(self.gamma, inputs.dtype) * inputs
//  [python]
def call(self, x, training=None):
    return x * t.cast(self._gamma, x.dtype)
//  [python]
def call(self, x, training=None):
    return x * tf.cast(self._gamma, x.dtype)
//  [python]
def call(self, inputs):
    """Run layer computation."""
    x, m = inputs

    reshaped_x = self._reshape_input(x)
    reshaped_m = self._reshape_input(m)

    q = tf.einsum('bnd,hkd->bnhk', reshaped_x, self._query_proj)
    k = tf.einsum('bmd,dk->bmk', reshaped_m, self._key_proj)
    logits = tf.einsum('bnhk,bmk->bnhm', q, k)

    logits = logits / tf.math.sqrt(tf.to_float(self._key_dim))
    attention_scores = self._dropout_layer(tf.nn.softmax(logits))

    v = tf.einsum('bmd,dv->bmv', reshaped_m, self._value_proj)
    o = tf.einsum('bnhm,bmv->bnhv', attention_scores, v)
    result = tf.einsum('bnhv,dhv->bnd', o, self._output_proj)

    return tf.ensure_shape(tf.reshape(result, tf.shape(x)), x.shape)
//  [python]
def call(self, inputs):
    """Run layer computation."""
    x, m = inputs

    reshaped_x = self._reshape_input(x)
    reshaped_m = self._reshape_input(m)

    q = tf.einsum('bnd,hkd->bnhk', reshaped_x, self._query_proj)
    k = tf.einsum('bmd,dk->bmk', reshaped_m, self._key_proj)
    logits = tf.einsum('bnhk,bmk->bnhm', q, k)

    logits = logits / tf.math.sqrt(tf.cast(self._key_dim, x.dtype))
    attention_scores = self._dropout_layer(tf.nn.softmax(logits))

    v = tf.einsum('bmd,dv->bmv', reshaped_m, self._value_proj)
    o = tf.einsum('bnhm,bmv->bnhv', attention_scores, v)
    result = tf.einsum('bnhv,dhv->bnd', o, self._output_proj)

    return tf.ensure_shape(tf.reshape(result, tf.shape(x)), x.shape)
//  [python]
def call(self, inputs, optimize_einsum=False):
    """Run layer computation."""
    x, m = inputs

    reshaped_x = self._reshape_input(x)
    reshaped_m = self._reshape_input(m)

    if optimize_einsum:
      logits = tf.einsum(
          'bnd,bme,hdk,ek->bhnm',
          reshaped_x,
          reshaped_m,
          self._query_proj,
          self._key_proj,
          optimize='optimal',
      )
    else:
      q = tf.einsum('bnd,hdk->bhnk', reshaped_x, self._query_proj)
      k = tf.einsum('bmd,dk->bmk', reshaped_m, self._key_proj)
      logits = tf.einsum('bhnk,bmk->bhnm', q, k)

    logits = logits / tf.math.sqrt(tf.to_float(self._key_dim))

    attention_scores = self._dropout_layer(tf.nn.softmax(logits))

    if optimize_einsum:
      result = tf.einsum(
          'bhnm,bmd,dv,hev->bne',
          attention_scores,
          reshaped_m,
          self._value_proj,
          self._output_proj,
          optimize='optimal',
      )
    else:
      v = tf.einsum('bmd,dv->bmv', reshaped_m, self._value_proj)
      o = tf.einsum('bhnm,bmv->bhnv', attention_scores, v)
      result = tf.einsum('bhnv,hdv->bnd', o, self._output_proj)

    return tf.ensure_shape(tf.reshape(result, tf.shape(x)), x.shape)
//  [python]
def call(self, inputs, optimize_einsum=False):
    """Run layer computation."""
    x, m = inputs

    reshaped_x = self._reshape_input(x)
    reshaped_m = self._reshape_input(m)

    if optimize_einsum:
      logits = tf.einsum(
          'bnd,bme,hdk,ek->bhnm',
          reshaped_x,
          reshaped_m,
          self._query_proj,
          self._key_proj,
          optimize='optimal',
      )
    else:
      q = tf.einsum('bnd,hdk->bhnk', reshaped_x, self._query_proj)
      k = tf.einsum('bmd,dk->bmk', reshaped_m, self._key_proj)
      logits = tf.einsum('bhnk,bmk->bhnm', q, k)

    logits = logits / tf.math.sqrt(tf.cast(self._key_dim, x.dtype))
    attention_scores = self._dropout_layer(tf.nn.softmax(logits))

    if optimize_einsum:
      result = tf.einsum(
          'bhnm,bmd,dv,hev->bne',
          attention_scores,
          reshaped_m,
          self._value_proj,
          self._output_proj,
          optimize='optimal',
      )
    else:
      v = tf.einsum('bmd,dv->bmv', reshaped_m, self._value_proj)
      o = tf.einsum('bhnm,bmv->bhnv', attention_scores, v)
      result = tf.einsum('bhnv,hdv->bnd', o, self._output_proj)

    return tf.ensure_shape(tf.reshape(result, tf.shape(x)), x.shape)
//  [python]
def create_scatter_indices_for_dim(dim, shape, indices=None):
  """Create scatter indieces for a given dimension."""
  dim_size = shape[dim]
  n_dims = len(shape)
  reshape = [1] * n_dims
  reshape[dim] = -1

  if indices is None:
    indices = tf.range(dim_size, dtype=tf.int32)  # [dim_size,]

  indices = tf.reshape(indices, reshape)  # [1, ..., dim_size, ..., 1]

  tf.debugging.assert_equal(tf.shape(indices)[dim], shape[dim])
  indices = tf.broadcast_to(
      indices, shape)  # [Ro, S, 1] or [Ro, S, C, 1]  [0,1,1,1] vs. [512,64,1,1]

  indices = tf.to_int32(indices)
  return indices
//  [python]
def create_scatter_indices_for_dim(dim, shape, indices=None):
  """Create scatter indieces for a given dimension."""
  dim_size = shape[dim]
  n_dims = len(shape)
  reshape = [1] * n_dims
  reshape[dim] = -1

  if indices is None:
    indices = tf.range(dim_size, dtype=tf.int32)  # [dim_size,]

  indices = tf.reshape(indices, reshape)  # [1, ..., dim_size, ..., 1]

  tf.debugging.assert_equal(tf.shape(indices)[dim], shape[dim])
  indices = tf.broadcast_to(
      indices, shape)  # [Ro, S, 1] or [Ro, S, C, 1]  [0,1,1,1] vs. [512,64,1,1]

  indices = tf.cast(indices, dtype=tf.int32)
  return indices
//  [python]
def _look_ahead_op(self, var):
    var_dtype = var.dtype.base_dtype
    slow_var = self.get_slot(var, 'slow')
    local_step = t.cast(self.iterations + 1, t.int64)
    sync_period = self._get_hyper('sync_period', t.int64)
    slow_step_size = self._get_hyper('slow_step_size', var_dtype)
    step_back = slow_var + slow_step_size * (var - slow_var)
    sync_cond = tf.equal(
        tf.math.floordiv(local_step, sync_period) * sync_period, local_step
    )
    with tf.control_dependencies([step_back]):
      slow_update = slow_var.assign(
          tf.where(sync_cond, step_back, slow_var),
          use_locking=self._use_locking,
      )
      var_update = var.assign(
          tf.where(sync_cond, step_back, var), use_locking=self._use_locking
      )
    return tf.group(slow_update, var_update)
//  [python]
def _look_ahead_op(self, var):
    var_dtype = var.dtype.base_dtype
    slow_var = self.get_slot(var, 'slow')
    local_step = tf.cast(self.iterations + 1, tf.dtypes.int64)
    sync_period = self._get_hyper('sync_period', tf.dtypes.int64)
    slow_step_size = self._get_hyper('slow_step_size', var_dtype)
    step_back = slow_var + slow_step_size * (var - slow_var)
    sync_cond = tf.equal(
        tf.math.floordiv(local_step, sync_period) * sync_period, local_step
    )
    with tf.control_dependencies([step_back]):
      slow_update = slow_var.assign(
          tf.where(sync_cond, step_back, slow_var),
          use_locking=self._use_locking,
      )
      var_update = var.assign(
          tf.where(sync_cond, step_back, var), use_locking=self._use_locking
      )
    return tf.group(slow_update, var_update)
//  [python]
def initialize(self):
    """Makes an initial state for the GaussianNoiseGenerator.

    Returns:
      A named tuple of (seeds, stddev).
    """
    if self._seed is None:
      time_now = tf.timestamp()
      residual = time_now - tf.math.floor(time_now)
      return self._GlobalState(
          tf.cast(
              tf.stack([
                  tf.math.floor(tf.timestamp() * 1e6),
                  tf.math.floor(residual * 1e9)
              ]),
              dtype=tf.int64), tf.constant(self._noise_std, dtype=tf.float32))
    else:
      return self._GlobalState(
          tf.constant(self._seed, dtype=tf.int64, shape=(2,)),
          tf.constant(self._noise_std, dtype=tf.float32))
//  [python]
def make_state(self, seeds: tf.Tensor, stddev: tf.Tensor):
    """Returns a new named tuple of (seeds, stddev)."""
    seeds = tf.ensure_shape(seeds, shape=(2,))
    return self._GlobalState(
        t.cast(seeds, dtype=tf.int64), t.cast(stddev, dtype=tf.float32))
//  [python]
def make_state(self, seeds: tf.Tensor, stddev: tf.Tensor):
    """Returns a new named tuple of (seeds, stddev)."""
    seeds = tf.ensure_shape(seeds, shape=(2,))
    return self._GlobalState(
        tf.cast(seeds, dtype=tf.int64), tf.cast(stddev, dtype=tf.float32))
//  [python]
def call(self, inputs):
    """Run layer computation."""
    x = inputs
    px = self._get_pixels(x)

    if self._query_h_strides > 1 or self._query_w_strides > 1:
      q = self._query_downsampling(x)
      q = self._query_downsampling_norm(q)
      q = self._query_proj(q)
    else:
      q = self._query_proj(x)

    # desired q shape: [b, n x n, h, k] - [b, l, h, k]
    q = self._reshape_projected_query(
        q,
        self._num_heads,
        px // self._query_h_strides,
        px // self._query_w_strides,
        self._key_dim,
    )

    if self._kv_strides > 1:
      k = self._key_dw_conv(x)
      k = self._key_dw_norm(k)
      k = self._key_proj(k)
    else:
      k = self._key_proj(x)
    # output shape of k: [b, k, p], p = m x m
    k = self._reshape_input(k)

    # desired q shape: [b, n x n, h, k]
    # desired k shape: [b, m x m, k]
    # desired logits shape: [b, n x n, h, m x m]
    logits = tf.einsum('blhk,bpk->blhp', q, k)

    logits = logits / tf.math.sqrt(tf.to_float(self._key_dim))

    attention_scores = self._dropout_layer(tf.nn.softmax(logits))

    if self._kv_strides > 1:
      v = self._value_dw_conv(x)
      v = self._value_dw_norm(v)
      v = self._value_proj(v)
    else:
      v = self._value_proj(x)

    # output shape of v: [ b, p, k], p = m x m
    v = self._reshape_input(v)
    o = tf.einsum('blhp,bpk->blhk', attention_scores, v)
    # reshape o into [b, n, n, hk]
    o = self._reshape_output(
        o,
        self._num_heads,
        px // self._query_h_strides,
        px // self._query_w_strides,
    )
    if self._query_h_strides > 1 or self._query_w_strides > 1:
      o = self._upsampling(o)

    result = self._output_proj(o)

    return tf.ensure_shape(tf.reshape(result, tf.shape(x)), x.shape)
//  [python]
def call(self, inputs):
    """Run layer computation."""
    x = inputs
    px = self._get_pixels(x)

    if self._query_h_strides > 1 or self._query_w_strides > 1:
      q = self._query_downsampling(x)
      q = self._query_downsampling_norm(q)
      q = self._query_proj(q)
    else:
      q = self._query_proj(x)

    # desired q shape: [b, n x n, h, k] - [b, l, h, k]
    q = self._reshape_projected_query(
        q,
        self._num_heads,
        px // self._query_h_strides,
        px // self._query_w_strides,
        self._key_dim,
    )

    if self._kv_strides > 1:
      k = self._key_dw_conv(x)
      k = self._key_dw_norm(k)
      k = self._key_proj(k)
    else:
      k = self._key_proj(x)
    # output shape of k: [b, k, p], p = m x m
    k = self._reshape_input(k)

    # desired q shape: [b, n x n, h, k]
    # desired k shape: [b, m x m, k]
    # desired logits shape: [b, n x n, h, m x m]
    logits = tf.einsum('blhk,bpk->blhp', q, k)

    logits = logits / tf.math.sqrt(tf.cast(self._key_dim, x.dtype))

    attention_scores = self._dropout_layer(tf.nn.softmax(logits))

    if self._kv_strides > 1:
      v = self._value_dw_conv(x)
      v = self._value_dw_norm(v)
      v = self._value_proj(v)
    else:
      v = self._value_proj(x)

    # output shape of v: [ b, p, k], p = m x m
    v = self._reshape_input(v)
    o = tf.einsum('blhp,bpk->blhk', attention_scores, v)
    # reshape o into [b, n, n, hk]
    o = self._reshape_output(
        o,
        self._num_heads,
        px // self._query_h_strides,
        px // self._query_w_strides,
    )
    if self._query_h_strides > 1 or self._query_w_strides > 1:
      o = self._upsampling(o)

    result = self._output_proj(o)

    return tf.ensure_shape(tf.reshape(result, tf.shape(x)), x.shape)
//  [python]
def _get_cumsum(self, state: TreeState) -> tf.Tensor:
    """Returns weighted cumulative sum of noise based on `TreeState`."""
    # Note that the buffer saved recursive results of the weighted average of
    # the node value (v) and its two children (l, r), i.e., node = v + (l+r)/2.
    # To get unbiased estimation with reduced variance for each node, we have to
    # reweight it by 1/(2-2^{-d}) where d is the depth of the node.
    level_weights = tf.math.divide(
        1., 2. - tf.math.pow(.5, tf.cast(state.level_buffer_idx, tf.float32)))

    def _weighted_sum(buffer):
      expand_shape = [len(level_weights)] + [1] * (len(tf.shape(buffer)) - 1)
      weighted_buffer = tf.math.multiply(
          buffer, tf.reshape(level_weights, expand_shape))
      return tf.reduce_sum(weighted_buffer, axis=0)

    return tf.nest.map_structure(_weighted_sum, state.level_buffer)
//  [python]
def gelu(features, approximate=False):
    if approximate:
        coeff = t.cast(0.044715, features.dtype)
        return 0.5 * features * (
            1.0 + tf.tanh(0.7978845608028654 *
                                (features + coeff * tf.pow(features, 3))))
    else:
        return 0.5 * features * (1.0 + tf.erf(
            features / t.cast(1.4142135623730951, features.dtype)))
//  [python]
def gelu(features, approximate=False):
    if approximate:
        coeff = tf.cast(0.044715, features.dtype)
        return 0.5 * features * (
            1.0 + tf.tanh(0.7978845608028654 *
                                (features + coeff * tf.pow(features, 3))))
    else:
        return 0.5 * features * (1.0 + tf.erf(
            features / tf.cast(1.4142135623730951, features.dtype)))
//  [python]
def get_angles(pos, i, dim):
    angle_rates = 1 / t.pow(10000., t.cast(2 * (i//2), t.float32) / dim)
    return t.cast(pos, t.float32) * t.cast(angle_rates, t.float32)
//  [python]
def get_angles(pos, i, dim):
    angle_rates = 1 / tf.pow(10000., tf.cast(2 * (i//2), tf.float32) / dim)
    return tf.cast(pos, tf.float32) * tf.cast(angle_rates, tf.float32)
//  [python]
def scatter_results(intersect, n_rays, keys):
  """Scatters intersecting ray results into the original set of rays.

  Args:
    intersect: Dict. Intersecting values.
    n_rays: int or tf.int32. Total number of rays.
    keys: [str]. List of keys to scatter.

  Returns:
    scattered_results: Dict. Scattered results.
  """
  # We use `None` to indicate that the intersecting set of rays is equivalent to
  # the full set of rays, so we are done.
  intersect_indices = intersect['indices']
  if intersect_indices is None:
    return {k: intersect[k] for k in keys}

  scattered_results = {}
  n_samples = intersect['z_vals'].shape[1]
  dim2known_indices = {0: intersect_indices}  # [R?, 1]
  for k in keys:
    if k == 'z_vals':
      tensor = tf.random.uniform((n_rays, n_samples),
                                 dtype=tf.float32)  # [R, S]
    elif k == 'pts':
      tensor = t.cast(  # [R, S, 3]
          tf.fill((n_rays, n_samples, 3), 1000.0),
          dtype=tf.float32)
    elif 'rgb' in k:
      tensor = tf.zeros((n_rays, n_samples, 3), dtype=tf.float32)  # [R, S, 3]
    elif 'alpha' in k:
      tensor = tf.zeros((n_rays, n_samples, 1), dtype=tf.float32)  # [R, S, 1]
    else:
      raise ValueError(f'Invalid key: {k}')
    scattered_v = scatter_nd(  # [R, S, K]
        tensor=tensor,
        updates=intersect[k],  # [Ro, S]
        dim2known_indices=dim2known_indices)
    # Convert the batch dimension to a known dimension.
    # For some reason `scattered_z_vals` becomes [R, ?]. We need to explicitly
    # reshape it with `n_samples`.
    if k == 'z_vals':
      scattered_v = tf.reshape(scattered_v, (n_rays, n_samples))  # [R, S]
    else:
      # scattered_v = tf.reshape(
      #     scattered_v, (n_rays,) + scattered_v.shape[1:])  # [R, S, K]
      # scattered_v = tf.reshape(
      # scattered_v, (-1,) + scattered_v.shape[1:])  # [R, S, K]
      scattered_v = tf.reshape(
          scattered_v, (n_rays, n_samples, tensor.shape[2]))  # [R, S, K]
    scattered_results[k] = scattered_v
  return scattered_results
//  [python]
def scatter_results(intersect, n_rays, keys):
  """Scatters intersecting ray results into the original set of rays.

  Args:
    intersect: Dict. Intersecting values.
    n_rays: int or tf.int32. Total number of rays.
    keys: [str]. List of keys to scatter.

  Returns:
    scattered_results: Dict. Scattered results.
  """
  # We use `None` to indicate that the intersecting set of rays is equivalent to
  # the full set of rays, so we are done.
  intersect_indices = intersect['indices']
  if intersect_indices is None:
    return {k: intersect[k] for k in keys}

  scattered_results = {}
  n_samples = intersect['z_vals'].shape[1]
  dim2known_indices = {0: intersect_indices}  # [R?, 1]
  for k in keys:
    if k == 'z_vals':
      tensor = tf.random.uniform((n_rays, n_samples),
                                 dtype=tf.float32)  # [R, S]
    elif k == 'pts':
      tensor = tf.cast(  # [R, S, 3]
          tf.fill((n_rays, n_samples, 3), 1000.0),
          dtype=tf.float32)
    elif 'rgb' in k:
      tensor = tf.zeros((n_rays, n_samples, 3), dtype=tf.float32)  # [R, S, 3]
    elif 'alpha' in k:
      tensor = tf.zeros((n_rays, n_samples, 1), dtype=tf.float32)  # [R, S, 1]
    else:
      raise ValueError(f'Invalid key: {k}')
    scattered_v = scatter_nd(  # [R, S, K]
        tensor=tensor,
        updates=intersect[k],  # [Ro, S]
        dim2known_indices=dim2known_indices)
    # Convert the batch dimension to a known dimension.
    # For some reason `scattered_z_vals` becomes [R, ?]. We need to explicitly
    # reshape it with `n_samples`.
    if k == 'z_vals':
      scattered_v = tf.reshape(scattered_v, (n_rays, n_samples))  # [R, S]
    else:
      # scattered_v = tf.reshape(
      #     scattered_v, (n_rays,) + scattered_v.shape[1:])  # [R, S, K]
      # scattered_v = tf.reshape(
      # scattered_v, (-1,) + scattered_v.shape[1:])  # [R, S, K]
      scattered_v = tf.reshape(
          scattered_v, (n_rays, n_samples, tensor.shape[2]))  # [R, S, K]
    scattered_results[k] = scattered_v
  return scattered_results
//  [python]
def positional_encoding(coords, dim):
    """coords in (bsz, size), return (bsz, size, dim)."""
    angle_rads = get_angles(tf.expand_dims(coords, -1),
                            tf.range(dim)[tf.newaxis, tf.newaxis, :],
                            dim)

    # apply sin to even indices in the array; 2i
    angle_rads1 = tf.sin(angle_rads[:, :, 0::2])

    # apply cos to odd indices in the array; 2i+1
    angle_rads2 = tf.cos(angle_rads[:, :, 1::2])

    pos_encoding = tf.concat([angle_rads1, angle_rads2], -1)

    return t.cast(pos_encoding, dtype=tf.float32)
//  [python]
def positional_encoding(coords, dim):
    """coords in (bsz, size), return (bsz, size, dim)."""
    angle_rads = get_angles(tf.expand_dims(coords, -1),
                            tf.range(dim)[tf.newaxis, tf.newaxis, :],
                            dim)

    # apply sin to even indices in the array; 2i
    angle_rads1 = tf.sin(angle_rads[:, :, 0::2])

    # apply cos to odd indices in the array; 2i+1
    angle_rads2 = tf.cos(angle_rads[:, :, 1::2])

    pos_encoding = tf.concat([angle_rads1, angle_rads2], -1)

    return tf.cast(pos_encoding, dtype=tf.float32)
//  [python]
def call(self, hidden_states, context, attention_mask, output_attentions, training=False):
        batch_size = shape_list(hidden_states)[0]
        mixed_query_layer = self.query(hidden_states)
        mixed_key_layer = self.key(context)
        mixed_value_layer = self.value(context)

        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = tf.matmul(
            query_layer, key_layer, transpose_b=True
        )  # (batch size, num_heads, seq_len_q, seq_len_k)
        dk = tf.cast(shape_list(key_layer)[-1], dtype=attention_scores.dtype)  # scale attention_scores
        attention_scores = attention_scores / tf.math.sqrt(dk)

        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in TFLxmertModel call() function)
            attention_mask = tf.cast(attention_mask, dtype=attention_scores.dtype)
            attention_scores = attention_scores + attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = stable_softmax(attention_scores, axis=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs, training=training)
        context_layer = tf.matmul(attention_probs, value_layer)

        context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])
        context_layer = tf.reshape(
            context_layer, (batch_size, -1, self.all_head_size)
        )  # (batch_size, seq_len_q, all_head_size)

        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
        return outputs
//  [python]
def prepare_inputs(self):
        """ Prepares Input
        """
        q1_inputs, self.qmax = clip_sentence(self.q1_inputs, self.q1_len)
        q2_inputs, self.a1max = clip_sentence(self.q2_inputs, self.q2_len)
        q3_inputs, self.a2max = clip_sentence(self.q3_inputs, self.q3_len)

        self.q1_mask = tf.cast(q1_inputs, tf.bool)
        self.q2_mask = tf.cast(q2_inputs, tf.bool)
        self.q3_mask = tf.cast(q3_inputs, tf.bool)

        with tf.device('/cpu:0'):
            q1_embed = tf.nn.embedding_lookup(self.embeddings,
                                                    q1_inputs)
            q2_embed = tf.nn.embedding_lookup(self.embeddings,
                                                    q2_inputs)
            q3_embed = tf.nn.embedding_lookup(self.embeddings,
                                                    q3_inputs)

        if(self.args.all_dropout):
            # By default, this is disabled
            q1_embed = tf.nn.dropout(q1_embed, self.emb_dropout)
            q2_embed = tf.nn.dropout(q2_embed, self.emb_dropout)
            q3_embed = tf.nn.dropout(q3_embed, self.emb_dropout)

        # Ignore these. :)
        self.c1_cnn, self.c2_cnn, self.c3_cnn = None, None, None
        self.p1_pos, self.p2_pos, self.p3_pos = None, None, None

        if('TNET' in self.args.rnn_type):
            t_inputs, _ = clip_sentence(self.trans_inputs, self.trans_len)
            self.trans_embed = tf.nn.embedding_lookup(self.embeddings,
                                                        t_inputs)

        if('EXP' in self.args.rnn_type):
            p_user_inputs1, self.pair_user_q1max = clip_sentence(self.pair_user_inputs1, self.pair_user_len1)
            p_user_inputs2, self.pair_user_q2max = clip_sentence(self.pair_user_inputs2, self.pair_user_len2)
            pair_user_q1_embed = tf.nn.embedding_lookup(self.embeddings,
                                                        p_user_inputs1)
            pair_user_q2_embed = tf.nn.embedding_lookup(self.embeddings,
                                                        p_user_inputs2)

            self.pair_user_q1_embed = tf.nn.dropout(pair_user_q1_embed, self.emb_dropout)
            self.pair_user_q2_embed = tf.nn.dropout(pair_user_q2_embed, self.emb_dropout)

            self.pair_user_q1_mask = tf.cast(p_user_inputs1, tf.bool)
            self.pair_user_q2_mask = tf.cast(p_user_inputs2, tf.bool)

        self.q1_embed = q1_embed
        self.q2_embed = q2_embed
        self.q3_embed = q3_embed
//  [python]
def get_2d_position_codes(height, width, out_dim, normalization_max=6.2831852):
    """Get 2d positional embedding with sin/cos codes.
    Args:
      height: a `int` specifying the height of the 2d image / feature map.
      width: a `int` specifying the width of the 2d image / feature map.
      out_dim: a `int` specifying the output dimension of the encoding.
        Must be divisible by 2.
      normalization_max: normalize coordinates between [0, normalization_max].
        If None, raw coordinates from 0 to height/width will be used.
    Returns:
      positional code of shape (1, height, width, out_dim)
    """
    y_coords = t.cast(t.range(height), t.float32)
    if normalization_max is not None:
        y_coords = y_coords / (t.cast(height, t.float32) - 1) * normalization_max
    y_coords = positional_encoding(y_coords, out_dim//2)
    y_coords = t.expand_dims(y_coords, 2)
    y_coords = t.concat([y_coords, t.zeros_like(y_coords)], -1)

    x_coords = t.cast(t.range(width), t.float32)
    if normalization_max is not None:
        x_coords = x_coords / (t.cast(width, t.float32) - 1) * normalization_max
    x_coords = positional_encoding(x_coords, out_dim//2)
    x_coords = t.expand_dims(x_coords, 1)
    x_coords = t.concat([t.zeros_like(x_coords), x_coords], -1)

    return y_coords + x_coords
//  [python]
def get_2d_position_codes(height, width, out_dim, normalization_max=6.2831852):
    """Get 2d positional embedding with sin/cos codes.
    Args:
      height: a `int` specifying the height of the 2d image / feature map.
      width: a `int` specifying the width of the 2d image / feature map.
      out_dim: a `int` specifying the output dimension of the encoding.
        Must be divisible by 2.
      normalization_max: normalize coordinates between [0, normalization_max].
        If None, raw coordinates from 0 to height/width will be used.
    Returns:
      positional code of shape (1, height, width, out_dim)
    """
    y_coords = tf.cast(tf.range(height), tf.float32)
    if normalization_max is not None:
        y_coords = y_coords / (tf.cast(height, tf.float32) - 1) * normalization_max
    y_coords = positional_encoding(y_coords, out_dim//2)
    y_coords = tf.expand_dims(y_coords, 2)
    y_coords = tf.concat([y_coords, tf.zeros_like(y_coords)], -1)

    x_coords = tf.cast(tf.range(width), tf.float32)
    if normalization_max is not None:
        x_coords = x_coords / (tf.cast(width, tf.float32) - 1) * normalization_max
    x_coords = positional_encoding(x_coords, out_dim//2)
    x_coords = tf.expand_dims(x_coords, 1)
    x_coords = tf.concat([tf.zeros_like(x_coords), x_coords], -1)

    return y_coords + x_coords
//  [python]
def make_dense(targets, noclass, low_param):
  """Move a batch of targets to a dense 1-hot representation."""
  low = low_param / float(noclass - 1)
  high = 1.0 - low * (noclass - 1)
  targets = t.cast(targets, t.int64)
  return tf.one_hot(targets, depth=noclass, on_value=high, off_value=low)
//  [python]
def make_dense(targets, noclass, low_param):
  """Move a batch of targets to a dense 1-hot representation."""
  low = low_param / float(noclass - 1)
  high = 1.0 - low * (noclass - 1)
  targets = tf.cast(targets, tf.int64)
  return tf.one_hot(targets, depth=noclass, on_value=high, off_value=low)
//  [python]
def go_for_image(self, faces, load_image_first: bool = False, detect_faces_first: bool = True):
		y_map = {0: "real", 1: "fake"}
		aaa1 = faces

		if load_image_first:
			try:
				faces = self.load_image(faces)
				if not detect_faces_first:
					faces = [faces]
			except:
				print("error")
				return

		if detect_faces_first:
			faces = self.faceExtractor.extract([faces])[0]
		
		for face in faces:
			try:
				face = tf.image.resize(face, (self.input_shape[0], self.input_shape[1]), method="nearest")
				aa = tf.nn.softmax(self.model(tf.expand_dims(tf.to_float(face)/255., 0)))
				if np.argmax(aa) == 0:
					print(y_map[np.argmax(aa)])
					print(aa[0][0]*100)
					print(aa[0][1]*100)
					print(aaa1)
					self.i += 1
					self.mistaken.append(aaa1)
					print("----------------------------------------")

				# cv2.imshow("face", face.numpy())
				# cv2.waitKey(0)
			except:
				print("error")
				continue
//  [python]
def go_for_image(self, faces, load_image_first: bool = False, detect_faces_first: bool = True):
		y_map = {0: "real", 1: "fake"}
		aaa1 = faces

		if load_image_first:
			try:
				faces = self.load_image(faces)
				if not detect_faces_first:
					faces = [faces]
			except:
				print("error")
				return

		if detect_faces_first:
			faces = self.faceExtractor.extract([faces])[0]
		
		for face in faces:
			try:
				face = tf.image.resize(face, (self.input_shape[0], self.input_shape[1]), method="nearest")
				aa = tf.nn.softmax(self.model(tf.expand_dims(tf.cast(face, tf.float32)/255., 0)))
				if np.argmax(aa) == 0:
					print(y_map[np.argmax(aa)])
					print(aa[0][0]*100)
					print(aa[0][1]*100)
					print(aaa1)
					self.i += 1
					self.mistaken.append(aaa1)
					print("----------------------------------------")

				# cv2.imshow("face", face.numpy())
				# cv2.waitKey(0)
			except:
				print("error")
				continue
//  [python]
def _process_label(label: tf.Tensor,
                   one_hot_label: bool = True,
                   num_classes: Optional[int] = None) -> tf.Tensor:
  """Processes label Tensor."""
  # Validate parameters.
  if one_hot_label and not num_classes:
    raise ValueError(
        '`num_classes` should be given when requesting one hot label.')

  # Cast to tf.int32.
  label = tf.cast(label, dtype=tf.int32)

  if one_hot_label:
    # Replace label index by one hot representation.
    label = tf.one_hot(label, num_classes)
    if len(label.shape.as_list()) > 1:
      label = tf.reduce_sum(label, axis=0)
    if num_classes == 1:
      # The trick for single label.
      label = 1 - label

  return label