import tensorflow as tf

        ],
    ]
    self.assertAllEqual(
        expected,
        relative_pos_gen.make_local_relative_att_ids(
            seq_len=3, local_radius=4, batch_size=2))

  def test_make_local_relative_att_ids_batch_size_2_tensor(self):
    dummy_batch = tf.ones([2, 5])

    relative_pos_gen = feature_utils.RelativePositionGenerator(max_distance=3)

    expected = [
        [
            [6, 6, 5, 4, 0, 1, 2, 3, 3],  #
            [6, 6, 5, 4, 0, 1, 2, 3, 3],  #
