import tensorflow as tf

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

  def test_make_local_relative_att_ids_invalid_arguments(self):
    relative_pos_gen = feature_utils.RelativePositionGenerator(max_distance=3)

    with self.assertRaises(ValueError):
      relative_pos_gen.make_local_relative_att_ids(seq_len=0, local_radius=3)

    with self.assertRaises(ValueError):
      relative_pos_gen.make_local_relative_att_ids(seq_len=5, local_radius=0)

    with self.assertRaises(ValueError):
