import tensorflow as tf

        relative_pos_gen.make_relative_att_ids(seq_len=5, batch_size=2))

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

    relative_pos_gen = feature_utils.RelativePositionGenerator(max_distance=3)
