import tensorflow as tf

                direct_mask_un = tf.greater(unhead_idxs, undep_idxs)  # [bs, sluh, sld]
            else:
                direct_mask_un = tf.less(unhead_idxs, undep_idxs)  # [bs, sluh, sld]

        # [bs, sluh, sld]
        rep_mask_tile_un = tf.logical_and(tf.expand_dims(rep_dep_mask, 1), tf.expand_dims(rep_unhead_mask, 2))
        pooling_mask = tf.logical_and(direct_mask_un, rep_mask_tile_un)  # [bs, sluh, sld]

        # data for pooling
        pooling_data = tf.tile(tf.expand_dims(rep_dep_tensor, 1), [1, sl_unhead, 1, 1])  # bs,sluh,sld,hn
