import tensorflow as tf

            [6, 5, 4, 0, 1],  #
            [6, 6, 5, 4, 0],  #
        ]
    ]
    self.assertAllEqual(
        expected,
        relative_pos_gen.make_relative_att_ids(
            seq_len=5, batch_size=tf.shape(dummy_batch)[0]))

  def test_overwrite_relative_att_ids_outside_segments(self):

    # batch_size = 2, seq_len = 5, max_distance = 3
    rel_att_ids = [
        [
