import tensorflow as tf

    with tf.name_scope('pooling_for_un_head'):
        undep_idxs = tf.tile(tf.expand_dims(dep_org_idx, 1), [1, sl_unhead, 1])  # [bs, sluh, sld]
        unhead_idxs = tf.tile(tf.expand_dims(unhead_org_idx, 2), [1, 1, sl_dep])  # [bs, sluh, sld]
        if direction is None:
