Compute the max first for tree decoding

[vasushyam]

Hi, I'm one of the authors of arXiv:2408.04093. I came to know from a mutual contact that you rightly pointed out that the version of Algorithm 3 we have in v2 of the paper is not strictly overflow safe because of how the maxes are handled. Here's a reworked version that should just be stable:

Thank you for pointing this out! and great work with the rapid implementation.

[lucidrains]

@vasushyam hey Vasudev! yes that does look good

thank you for uncovering the neat algorithm!