Question about topk_index = topk_index.view(-1) in cnets.py

[HaoWuSR]

Dear Authors, Thank you for your outstanding work on EAGLE and EAGLE-2! I have a question regarding the implementation, specifically about the line topk_index = topk_index.view(-1). Let me illustrate my confusion with an example: When executing topk_index, topk_prob, op = self.sample(last_headout, logits_processor, k=top_k), with top_k=3, we obtain 4 branches: [ [55, 67, 33], [25578, 13, 45], [2, 56, 42], [32, 64, 89] ] After applying topk_index = topk_index.view(-1), the output becomes: [55, 67, 33, 25578, 13, 45, 2, 56, 42, 32, 64, 89]

Then, if self.tree_buffer['tree_indices'][i] = [0, 2], we get select_index = [55, 33]. Is my understanding correct? Initially, I expected the implementation to select tokens from each branch. However, it appears we're selecting tokens from specific positions in the flattened array. Is there a particular reason for ignoring certain branches? Or is this a deliberate design choice in the tree buffer indices? I would greatly appreciate ur time! Regards,