1 / 1

Pad and Mask to Bucket Length

JAX

Medium

Turn a ragged batch of token sequences, stored back to back in one flat array, into a padded batch and its mask, with the padding work compiled by jax.jit once per bucket length.

tokens is a flat int32 JAX array that holds the B sequences of the batch one after another, and lengths is an int32 JAX array of shape (B,) whose entry lengths[b] is the number of tokens in sequence b. Writing $\ell_b$ for lengths[b], sequence b starts at the flat position

s_b = \sum_{c=0}^{b-1} \ell_c ,

and occupies positions $s_b$ to $s_b + \ell_b - 1$ .

buckets is a Python list of positive ints in any order, and pad_id is a Python int. The function is called as pad_to_bucket(tokens, lengths, buckets, pad_id).

The padded length L is the smallest value in buckets that is at least the longest length. In the output, row b and column j hold tokens[ $s_b + j$ ] when $j < \ell_b$ and pad_id otherwise, and the mask is True exactly where $j < \ell_b$ .

Return a dict with exactly two keys. tokens is an int32 JAX array of shape (B, L), and mask is a bool JAX array of shape (B, L).

Examples

Input: tokens = [5,6,7,8,9], lengths = [3,2], buckets = [4,8], pad_id = 0

Output: {"tokens":[[5,6,7,0],[8,9,0,0]],"mask":[[true,true,true,false],[true,true,false,false]]}

Explanation: The lengths add up to the 5 flat tokens, so the rows start at offsets 0 and 3. The longest sequence has 3 tokens and 4 is the smallest bucket that fits it, so both rows are padded with 0 up to length 4.

Input: tokens = [1,2,3,4,5,6], lengths = [5,1], buckets = [16,4,8], pad_id = -1

Output: {"tokens":[[1,2,3,4,5,-1,-1,-1],[6,-1,-1,-1,-1,-1,-1,-1]],"mask":[[true,true,true,true,true,false,false,false],[true,false,false,false,false,false,false,false]]}

Input: tokens = [0,3,0,0,2], lengths = [4,0,1], buckets = [2,4,6], pad_id = 0

Output: {"tokens":[[0,3,0,0],[0,0,0,0],[2,0,0,0]],"mask":[[true,true,true,true],[false,false,false,false],[true,false,false,false]]}

Hint 1

jax.jit(f, static_argnames="length") hands length to f as a Python int, so jnp.arange(length) has a fixed size.

Hint 2

jnp.cumsum(lengths) returns inclusive running totals, and a row starts where the earlier rows' total ends.

Hint 3

tokens[index] with a (B, L) integer array gathers all slots at once, and jnp.clip bounds index.

Requirements

Choose L outside jit from int(jnp.max(lengths)) and the buckets list, then pass it to a function compiled with jax.jit as a static argument through static_argnames or static_argnums.
Make the bucket length the only static value that depends on the batch, and define the function that jax.jit compiles once at module level. Two batches with the same array shapes in the same bucket must then reuse one compiled program, while each new bucket compiles its own.
Inside the jitted function, compute the start offsets from lengths with jnp.cumsum, build the positions with jnp.arange(L), read all B by L slots from the flat tokens array in one gather such as tokens[index] with a (B, L) index array, and write pad_id into the padded slots with jnp.where. Indices past the end of tokens only occur in padded slots, so clipping them with jnp.clip is optional.
Do not loop over rows in Python, inside jit or outside it, and do not convert tokens to NumPy or to a Python list. The grader traces the jitted function again for larger batches and expects the same number of operations, which a loop over rows breaks. The lengths are read on the host only to find the longest length.
Build the mask from lengths, because a real token can have the same id as pad_id.
Leave the buckets list unchanged.

Constraints

lengths holds between 1 and 6 non-negative integers, each at most 32, and their sum is at least 1.
tokens holds exactly as many int32 token ids as the sum of lengths, and every id lies between 0 and 59,999.
buckets holds between 1 and 5 distinct positive integers, each at most 64, in any order, and at least one of them is at least the longest length.
pad_id is an integer between -1 and 99,999.

Try Similar Problems

Functional Array UpdateEasy PRNG Key SplittingEasy Masked Mean Without Boolean IndexingEasy 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 Global Norm Gradient ClippingMedium

Pad and Mask to Bucket Length

JAX

Medium

Turn a ragged batch of token sequences, stored back to back in one flat array, into a padded batch and its mask, with the padding work compiled by jax.jit once per bucket length.

s_b = \sum_{c=0}^{b-1} \ell_c ,

and occupies positions $s_b$ to $s_b + \ell_b - 1$ .

buckets is a Python list of positive ints in any order, and pad_id is a Python int. The function is called as pad_to_bucket(tokens, lengths, buckets, pad_id).

Return a dict with exactly two keys. tokens is an int32 JAX array of shape (B, L), and mask is a bool JAX array of shape (B, L).

Examples

Input: tokens = [5,6,7,8,9], lengths = [3,2], buckets = [4,8], pad_id = 0

Output: {"tokens":[[5,6,7,0],[8,9,0,0]],"mask":[[true,true,true,false],[true,true,false,false]]}

Input: tokens = [1,2,3,4,5,6], lengths = [5,1], buckets = [16,4,8], pad_id = -1

Output: {"tokens":[[1,2,3,4,5,-1,-1,-1],[6,-1,-1,-1,-1,-1,-1,-1]],"mask":[[true,true,true,true,true,false,false,false],[true,false,false,false,false,false,false,false]]}

Input: tokens = [0,3,0,0,2], lengths = [4,0,1], buckets = [2,4,6], pad_id = 0

Output: {"tokens":[[0,3,0,0],[0,0,0,0],[2,0,0,0]],"mask":[[true,true,true,true],[false,false,false,false],[true,false,false,false]]}

Hint 1

jax.jit(f, static_argnames="length") hands length to f as a Python int, so jnp.arange(length) has a fixed size.

Hint 2

jnp.cumsum(lengths) returns inclusive running totals, and a row starts where the earlier rows' total ends.

Hint 3

tokens[index] with a (B, L) integer array gathers all slots at once, and jnp.clip bounds index.

Requirements

Choose L outside jit from int(jnp.max(lengths)) and the buckets list, then pass it to a function compiled with jax.jit as a static argument through static_argnames or static_argnums.
Make the bucket length the only static value that depends on the batch, and define the function that jax.jit compiles once at module level. Two batches with the same array shapes in the same bucket must then reuse one compiled program, while each new bucket compiles its own.
Inside the jitted function, compute the start offsets from lengths with jnp.cumsum, build the positions with jnp.arange(L), read all B by L slots from the flat tokens array in one gather such as tokens[index] with a (B, L) index array, and write pad_id into the padded slots with jnp.where. Indices past the end of tokens only occur in padded slots, so clipping them with jnp.clip is optional.
Do not loop over rows in Python, inside jit or outside it, and do not convert tokens to NumPy or to a Python list. The grader traces the jitted function again for larger batches and expects the same number of operations, which a loop over rows breaks. The lengths are read on the host only to find the longest length.
Build the mask from lengths, because a real token can have the same id as pad_id.
Leave the buckets list unchanged.

Constraints

lengths holds between 1 and 6 non-negative integers, each at most 32, and their sum is at least 1.
tokens holds exactly as many int32 token ids as the sum of lengths, and every id lies between 0 and 59,999.
buckets holds between 1 and 5 distinct positive integers, each at most 64, in any order, and at least one of them is at least the longest length.
pad_id is an integer between -1 and 99,999.

Try Similar Problems

tokens =

Accepts: array

lengths =

Accepts: array

buckets =

Accepts: array

pad_id =

Accepts: integer