1 / 1

Greedy Decode Loop with Scan

JAX

Medium

Generate num_steps tokens greedily from a tiny recurrent language model, running the decoding loop with jax.lax.scan.

The model keeps one hidden state per sequence. At step $t$ it reads the previous token $y_{t-1}$ and the previous hidden state $h_{t-1}$ , updates the state, scores every vocabulary entry, and picks the highest-scoring token.

h_t = \tanh\left(E_{y_{t-1}} + h_{t-1} W_{hh} + b_h\right), \qquad \ell_t = h_t W_{out} + b_{out}

y_t = \arg\max_{v} \ell_{t,v}, \qquad r_t = \log \operatorname{softmax}(\ell_t)_{y_t}

Here params is a dict holding the embedding table embed $E$ of shape (V, H), whose row $E_{y}$ is the vector for token $y$ , the recurrent weights w_hh $W_{hh}$ of shape (H, H), the bias b_h $b_h$ of shape (H,), the output weights w_out $W_{out}$ of shape (H, V) and the output bias b_out $b_{out}$ of shape (V,). V is the vocabulary size and H is the hidden size. The loop starts from $y_0$ = start_tokens and $h_0$ = h0, and steps $t = 1, \dots, T$ with $T$ = num_steps. $\ell_t$ holds the V logits of step $t$ , $y_t$ is the generated token and $r_t$ is its log probability. Every one of the B sequences in the batch runs the same rule independently.

Return a dict with two keys. tokens is an int32 JAX array of shape (B, num_steps) whose row $b$ lists $y_1, \dots, y_T$ for sequence $b$ , and log_probs is a float32 JAX array of shape (B, num_steps) holding the matching $r_t$ values.

Examples

Input: params = {"embed":[[-0.3,-0.8],[1.5,0.0],[-0.8,0.2],[1.7,1.5]],"w_hh":[[-0.6,0.4],[-0.1,-0.2]],"b_h":[0.1,0.1],"w_out":[[-1.6,-0.3,3.2,-2.3],[-0.4,0.9,-0.3,-2.0]],"b_out":[0.4,-0.3,-0.3,0.2]}, start_tokens = [0], h0 = [[0.0,0.0]], num_steps = 3

Output: {"tokens":[[3,2,0]],"log_probs":[[-0.43737,-0.117443,-0.703083]]}

Explanation: The first step reads token 0 with a zero state, so the new state is tanh(E_0 + b_h) = tanh([-0.2, -0.7]), about [-0.1974, -0.6044]. Its logits are about [0.9575, -0.7847, -0.7503, 1.8627], so the first token is 3 with log probability -0.43737. Feeding token 3 and the new state back into the same step gives token 2 and then token 0.

Input: params = {"embed":[[1.7,1.4,-1.1],[-0.1,-1.1,-1.1],[0.9,-0.8,0.1],[0.7,-0.7,-0.9],[0.2,0.5,-0.9]],"w_hh":[[0.2,-0.1,0.7],[0.0,0.5,-0.5],[0.0,0.2,-1.0]],"b_h":[0.5,0.1,0.3],"w_out":[[0.9,-1.0,1.5,-1.1,2.0],[-1.4,0.1,-0.9,-3.0,0.7],[-0.3,-1.0,-1.0,2.1,-3.7]],"b_out":[0.2,0.6,0.0,0.0,-0.1]}, start_tokens = [4,0], h0 = [[0.1,0.6,0.4],[-0.3,-0.2,0.9]], num_steps = 4

Output: {"tokens":[[4,4,4,4],[4,2,2,3]],"log_probs":[[-0.05271,-0.938562,-0.089923,-0.675244],[-0.019317,-1.258601,-1.079186,-0.434214]]}

Input: params = {"embed":[[0.8,-0.5],[-0.6,-1.0],[1.7,0.6],[-0.7,0.6]],"w_hh":[[-1.5,-0.9],[1.6,-0.3]],"b_h":[-0.2,-0.4],"w_out":[[0.2,-0.5,1.7,-0.4],[-2.0,0.7,-0.3,-0.8]],"b_out":[1.1,-0.8,0.4,0.5]}, start_tokens = [2,0,3], h0 = [[0.2,-0.4],[-0.2,-0.3],[-0.6,-0.1]], num_steps = 1

Output: {"tokens":[[2],[0],[3]],"log_probs":[[-0.792978],[-0.467053],[-1.220548]]}

Hint 1

jax.lax.scan(step, init, xs=None, length=n) runs step n times.

Hint 2

Carry (hidden, token); return (new_carry, (token, logprob)) to stack outputs.

Hint 3

jnp.take_along_axis(logp, tok[:, None], axis=-1) picks the chosen token's value.

Requirements

Run the loop over decoding steps with jax.lax.scan. The function is traced with jax.make_jaxpr, and the traced program must contain a scan, so a Python for loop that repeats the step num_steps times is rejected.
The function is called inside jax.jit with num_steps fixed as a Python int and every array traced, so it must not convert traced values to Python numbers.
Use float32 arithmetic for the hidden state, the logits and the log probabilities.

Constraints

B is between 1 and 4, V is between 4 and 8, and H is between 2 and 4.
num_steps is an integer between 1 and 8.
start_tokens is an int32 array of shape (B,) with values from 0 to V minus 1.
h0 is a float32 array of shape (B, H).
Every weight, bias and h0 value is a finite float32 number.
At every step the largest logit of every sequence is unique, so the greedy choice never ties.

Try Similar Problems

Paged KV Cache GatherMedium Beam Search StepHard Speculative Decoding AcceptanceHard Pad and Mask to Bucket LengthMedium Hessian-Vector ProductMedium Per-Example Gradient ClippingMedium Straight-Through Estimator with Custom VJPMedium Traced Learning Rate ScheduleMedium Exponential Moving Average with ScanMedium Spectral Norm by Power IterationMedium Batched Attention with vmapMedium Pairwise Distances with Nested vmapMedium