flax

4ccd5fff - Fix temperature sampler loop termination condition:

Fix temperature sampler loop termination condition: The `i0` starts from 0, so if we use `i <= max_decode_length`, we will run max_decoed_length + 1 times. PiperOrigin-RevId: 395652813