if valid_lens is not None:
    valid_lens = np.repeat(valid_lens, self.num_heads, axis=0)