//  [python]
def _diffpool(self, X, A, node_score, mask):

        hparams = self.hparams
        with tf.name_scope('diffpool'):
            k = 30
            for _ in range(self.pool_layers):
                # Update node embeddings
                Z = tf.layers.dense(X, 32, use_bias=False)

                # Compute cluster assignment matrix
                S = tf.layers.dense(X, k, use_bias=False) # B*L*F -> B*L*k
                S = tf.matmul(A, S) # B*L*L x B*L*k -> B*L*k
 
 
                # 2.via ratio
                num_nodes = tf.cast(tf.reduce_sum(mask, 1), tf.float32) # B
                node_position = tf.cast(tf.math.cumsum(mask, axis=1), tf.float32) # B*L 
                boolean_pool = tf.less(node_position, tf.expand_dims(self.pool_ratio*num_nodes, -1)) # B*L + B*1 -> B*L(k)
                # if k not 50
                boolean_pool = tf.batch_gather(boolean_pool, tf.tile(tf.expand_dims(tf.range(k), 0), [tf.shape(mask)[0], 1]))

                mask_paddings = tf.ones_like(S) * (-(2 ** 32) + 1) # B*L*L(k)
                S = tf.nn.softmax(
                    tf.where(tf.tile(tf.expand_dims(boolean_pool, 1), [1, tf.shape(mask)[1], 1]), S, mask_paddings),
                    axis = -1
                ) # B*1*L(k) + B*L*L(k) + B*L*(k) -> B*L*L(k)

                mask = tf.cast(boolean_pool, tf.int32) # B*L(k)

                ## Auxiliary pooling loss
                auxiliary_loss = 0.0
                # Link prediction loss
                S_gram = tf.matmul(S, tf.transpose(S, (0,2,1)))
                LP_loss = A - S_gram #  LP_loss = A/tf.norm(A) - S_gram/tf.norm(S_gram)
                LP_loss = tf.norm(LP_loss, axis=(-1, -2))
                LP_loss = K.mean(LP_loss)
                auxiliary_loss += self.same_mapping_regu*LP_loss
                # Entropy loss
                entr = tf.negative(tf.reduce_sum(tf.multiply(S, K.log(S + K.epsilon())), axis=-1))
                entr_loss = K.mean(entr, axis=-1)
                entr_loss = K.mean(entr_loss)
                auxiliary_loss += self.single_affiliation_regu*entr_loss
                # Position loss
                Pn = tf.math.cumsum(tf.ones([tf.shape(S)[0], 1, tf.shape(S)[1]]), axis=-1) # node position encoding:
                Pc = tf.math.cumsum(tf.ones([tf.shape(S)[0], 1, tf.shape(S)[2]]), axis=-1) # cluster position encoding:
                position_loss = tf.matmul(Pn, S) - Pc
                position_loss = tf.norm(position_loss, axis=(-1, -2))
                position_loss = K.mean(position_loss)
                auxiliary_loss += self.relative_position_regu*position_loss

                # pooled
                X = tf.matmul(tf.transpose(S, (0,2,1)), Z) # B*L(k)*L x B*L*F -> B*L(k)*F
                cluster_score = tf.squeeze(tf.matmul(tf.transpose(S, (0,2,1)), tf.expand_dims(node_score, -1)), -1) # B*L(k)*L x B*L*1 -> B*L(k)*1 -> B*L(k)


                X = X*tf.expand_dims(cluster_score, -1)
                A = tf.matmul(
                    tf.matmul(
                        tf.transpose(S, (0,2,1)), 
                        A), 
                    S) # B*L(k)*L x B*L*L x B*L*L(k) -> B*L(k)*L(k)


        return X, A, cluster_score, mask, auxiliary_loss