Perform one step of beam search for a batch by extending every beam with every vocabulary token, keeping the K best candidates and reordering the beam histories and caches so that each new beam carries the state of the beam it extends.
The batch has B rows, and each batch row b keeps K beams. Beam k has a running score beam_scores[b, k], the sum of the log probabilities of its tokens so far, and log_probs[b, k, v] is the log probability of token v coming next on that beam. Extending beam k by token v gives a candidate with total score
cb,k,v=beam_scores[b,k]+log_probs[b,k,v]Among all K times V candidates of row b, keep the K with the largest totals, sorted from best to worst. For new beam i, parents[b, i] is the beam k it extends, tokens[b, i] is the token v it adds, and scores[b, i] is its total cb,k,v. Its history is the history of beam parents[b, i] followed by tokens[b, i]. Here V is the vocabulary size and L is the number of tokens each history already holds.
cache is a dict of arrays that each start with the axes (B, K), one entry per beam, such as the attention keys and values of every beam. Every array in it is reordered the same way, so new beam i of row b holds the entry of beam parents[b, i].
Return a dict with five keys. scores is a float32 JAX array of shape (B, K), tokens and parents are int32 JAX arrays of shape (B, K), histories is an int32 JAX array of shape (B, K, L + 1), and cache is a dict with the same keys as the input cache whose arrays keep their shapes and dtypes.
Input: beam_scores = [[-1.1,-2.6]], log_probs = [[[-2.2,-1.7,-0.6,-1.8],[-2.4,-1.9,-3.0,-0.3]]], histories = [[[0,2],[0,1]]], cache = {"k":[[[-0.8,0.8],[0.8,0.3]]],"v":[[[0.1,0.0],[0.0,-0.3]]]}
Output: {"scores":[[-1.7,-2.8]],"tokens":[[2,1]],"parents":[[0,0]],"histories":[[[0,2,2],[0,2,1]]],"cache":{"k":[[[-0.8,0.8],[-0.8,0.8]]],"v":[[[0.1,0.0],[0.1,0.0]]]}}
Explanation: Beam 0 has score -1.1 and beam 1 has score -2.6. The best candidate is beam 0 with token 2, at -1.1 - 0.6 = -1.7. The second best is beam 0 with token 1, at -2.8, which beats the best candidate of beam 1 (token 3, at -2.9). Both new beams extend beam 0, so both copy its history [0, 2] and its cache rows, and the old beam 1 is dropped.
Input: beam_scores = [[-0.2,-0.7,-1.6],[-0.4,-1.4,-2.8]], log_probs = [[[-4.3,-4.7,-5.0,0.0],[-1.3,-2.5,-0.6,-2.1],[-3.5,-0.5,-1.2,-2.6]],[[-2.9,-0.3,-2.0,-2.5],[-0.6,-3.9,-1.1,-2.5],[-2.9,-1.2,-2.4,-0.6]]], histories = [[[0],[1],[0]],[[0],[2],[3]]], cache = {"k":[[[-0.4,-0.9],[0.8,-0.3],[-0.8,-0.5]],[[0.6,0.4],[0.7,-0.4],[-0.5,1.0]]],"v":[[[-0.4,-0.8],[0.9,0.0],[0.3,-1.0]],[[0.8,-0.1],[0.5,0.7],[-0.3,0.0]]]}
Output: {"scores":[[-0.2,-1.3,-2.0],[-0.7,-2.0,-2.4]],"tokens":[[3,2,0],[1,0,2]],"parents":[[0,1,1],[0,1,0]],"histories":[[[0,3],[1,2],[1,0]],[[0,1],[2,0],[0,2]]],"cache":{"k":[[[-0.4,-0.9],[0.8,-0.3],[0.8,-0.3]],[[0.6,0.4],[0.7,-0.4],[0.6,0.4]]],"v":[[[-0.4,-0.8],[0.9,0.0],[0.9,0.0]],[[0.8,-0.1],[0.5,0.7],[0.8,-0.1]]]}}
Input: beam_scores = [[0.0,-1000000000.0,-1000000000.0]], log_probs = [[[-4.9,-2.5,-2.2,-0.3,-3.0],[-1.9,-1.5,-0.7,-2.1,-3.5],[-2.0,-0.9,-3.7,-1.0,-3.2]]], histories = [[[],[],[]]], cache = {"k":[[[-0.2,-0.7],[-0.3,-0.1],[-0.6,-0.7]]],"v":[[[-1.0,-0.3],[-1.0,-0.1],[0.4,-0.4]]]}
Output: {"scores":[[-0.3,-2.2,-2.5]],"tokens":[[3,2,1]],"parents":[[0,0,0]],"histories":[[[3],[2],[1]]],"cache":{"k":[[[-0.2,-0.7],[-0.2,-0.7],[-0.2,-0.7]]],"v":[[[-1.0,-0.3],[-1.0,-0.3],[-1.0,-0.3]]]}}
jax.lax.top_k(totals.reshape(B, K * V), K) returns sorted scores and flat indices.
flat // V gives the parent beam and flat % V gives the token.
jax.vmap(lambda rows, picks: rows[picks])(leaf, parents) reorders one array, and jax.tree.map applies it to the cache.
Accepts: array
Accepts: array
Accepts: array
Accepts: object
Perform one step of beam search for a batch by extending every beam with every vocabulary token, keeping the K best candidates and reordering the beam histories and caches so that each new beam carries the state of the beam it extends.
The batch has B rows, and each batch row b keeps K beams. Beam k has a running score beam_scores[b, k], the sum of the log probabilities of its tokens so far, and log_probs[b, k, v] is the log probability of token v coming next on that beam. Extending beam k by token v gives a candidate with total score
cb,k,v=beam_scores[b,k]+log_probs[b,k,v]Among all K times V candidates of row b, keep the K with the largest totals, sorted from best to worst. For new beam i, parents[b, i] is the beam k it extends, tokens[b, i] is the token v it adds, and scores[b, i] is its total cb,k,v. Its history is the history of beam parents[b, i] followed by tokens[b, i]. Here V is the vocabulary size and L is the number of tokens each history already holds.
cache is a dict of arrays that each start with the axes (B, K), one entry per beam, such as the attention keys and values of every beam. Every array in it is reordered the same way, so new beam i of row b holds the entry of beam parents[b, i].
Return a dict with five keys. scores is a float32 JAX array of shape (B, K), tokens and parents are int32 JAX arrays of shape (B, K), histories is an int32 JAX array of shape (B, K, L + 1), and cache is a dict with the same keys as the input cache whose arrays keep their shapes and dtypes.
Input: beam_scores = [[-1.1,-2.6]], log_probs = [[[-2.2,-1.7,-0.6,-1.8],[-2.4,-1.9,-3.0,-0.3]]], histories = [[[0,2],[0,1]]], cache = {"k":[[[-0.8,0.8],[0.8,0.3]]],"v":[[[0.1,0.0],[0.0,-0.3]]]}
Output: {"scores":[[-1.7,-2.8]],"tokens":[[2,1]],"parents":[[0,0]],"histories":[[[0,2,2],[0,2,1]]],"cache":{"k":[[[-0.8,0.8],[-0.8,0.8]]],"v":[[[0.1,0.0],[0.1,0.0]]]}}
Explanation: Beam 0 has score -1.1 and beam 1 has score -2.6. The best candidate is beam 0 with token 2, at -1.1 - 0.6 = -1.7. The second best is beam 0 with token 1, at -2.8, which beats the best candidate of beam 1 (token 3, at -2.9). Both new beams extend beam 0, so both copy its history [0, 2] and its cache rows, and the old beam 1 is dropped.
Input: beam_scores = [[-0.2,-0.7,-1.6],[-0.4,-1.4,-2.8]], log_probs = [[[-4.3,-4.7,-5.0,0.0],[-1.3,-2.5,-0.6,-2.1],[-3.5,-0.5,-1.2,-2.6]],[[-2.9,-0.3,-2.0,-2.5],[-0.6,-3.9,-1.1,-2.5],[-2.9,-1.2,-2.4,-0.6]]], histories = [[[0],[1],[0]],[[0],[2],[3]]], cache = {"k":[[[-0.4,-0.9],[0.8,-0.3],[-0.8,-0.5]],[[0.6,0.4],[0.7,-0.4],[-0.5,1.0]]],"v":[[[-0.4,-0.8],[0.9,0.0],[0.3,-1.0]],[[0.8,-0.1],[0.5,0.7],[-0.3,0.0]]]}
Output: {"scores":[[-0.2,-1.3,-2.0],[-0.7,-2.0,-2.4]],"tokens":[[3,2,0],[1,0,2]],"parents":[[0,1,1],[0,1,0]],"histories":[[[0,3],[1,2],[1,0]],[[0,1],[2,0],[0,2]]],"cache":{"k":[[[-0.4,-0.9],[0.8,-0.3],[0.8,-0.3]],[[0.6,0.4],[0.7,-0.4],[0.6,0.4]]],"v":[[[-0.4,-0.8],[0.9,0.0],[0.9,0.0]],[[0.8,-0.1],[0.5,0.7],[0.8,-0.1]]]}}
Input: beam_scores = [[0.0,-1000000000.0,-1000000000.0]], log_probs = [[[-4.9,-2.5,-2.2,-0.3,-3.0],[-1.9,-1.5,-0.7,-2.1,-3.5],[-2.0,-0.9,-3.7,-1.0,-3.2]]], histories = [[[],[],[]]], cache = {"k":[[[-0.2,-0.7],[-0.3,-0.1],[-0.6,-0.7]]],"v":[[[-1.0,-0.3],[-1.0,-0.1],[0.4,-0.4]]]}
Output: {"scores":[[-0.3,-2.2,-2.5]],"tokens":[[3,2,1]],"parents":[[0,0,0]],"histories":[[[3],[2],[1]]],"cache":{"k":[[[-0.2,-0.7],[-0.2,-0.7],[-0.2,-0.7]]],"v":[[[-1.0,-0.3],[-1.0,-0.3],[-1.0,-0.3]]]}}
jax.lax.top_k(totals.reshape(B, K * V), K) returns sorted scores and flat indices.
flat // V gives the parent beam and flat % V gives the token.
jax.vmap(lambda rows, picks: rows[picks])(leaf, parents) reorders one array, and jax.tree.map applies it to the cache.
Accepts: array
Accepts: array
Accepts: array
Accepts: object
