Choose one token for each row of logits using temperature scaling, a top-k filter, a top-p (nucleus) filter and inverse-CDF sampling with the supplied uniform numbers.

For every row, apply these rules in order.

1. Divide the logits by temperature.
2. Keep the top_k tokens with the largest scaled logits and give every other token zero probability.
3. Take the softmax over the kept tokens to get their probabilities.
4. Rank the kept tokens from most to least likely. A token stays in the nucleus when the total probability of the tokens ranked before it is less than top_p, so the most likely token always stays and so does the token that first carries the total past top_p.
5. Divide the nucleus probabilities by their sum and give every token outside the nucleus probability 0.
6. Walk the vocabulary in index order 0, 1, ..., V - 1, keeping a running total $c_j$ of the final probabilities, and choose the smallest index $j$ with $c_j > u$, where $u$ is the row's entry of uniform. The comparison is strict, so a token with probability 0 is never chosen, even when $u$ is 0.

Here logits has shape (B, V) for B rows over a vocabulary of V tokens, uniform has shape (B,) with one number per row, temperature is the number the logits are divided by, top_k is how many tokens the top-k filter keeps, and top_p is the nucleus threshold.

Your function is called under jax.jit with logits and uniform traced and with temperature, top_k and top_p passed as static Python numbers, so every intermediate array needs a shape that does not depend on the data.

Return an int32 JAX array of shape (B,) holding the chosen token index for each row.