Study Plan

Static Arguments in jit

JAX

Easy

Normalize every attention head of a batch of token vectors on its own, inside a function compiled with jax.jit.

The input x has shape (B, T, D), where B is the batch size, T is the number of tokens and D is the feature width. Let H be num_heads and $D_h = D / H$ the head width. Head h of a token is the block of features from $h D_h$ to $(h + 1) D_h - 1$ . Each head vector $v$ with entries $v_1, \dots, v_{D_h}$ is normalized as

y = \frac{v}{\sqrt{\frac{1}{D_h} \sum_{j=1}^{D_h} v_j^2 + 10^{-6}}}

and the normalized heads are then merged back into vectors of width D in the original feature order.

Return a dict with the key heads, a float32 JAX array of shape (B, H, T, $D_h$ ) where entry [b, h, t] is the normalized head h of token t in sample b, and the key merged, a float32 JAX array of shape (B, T, D).

Examples

Input: x = [[[1,2,3,4],[0,0,3,-4]]], num_heads = 2

Output: {"heads":[[[[0.632455,1.264911],[0.0,0.0]],[[0.848528,1.131371],[0.848528,-1.131371]]]],"merged":[[[0.632455,1.264911,0.848528,1.131371],[0.0,0.0,0.848528,-1.131371]]]}

Explanation: Head 0 of the first token is [1,2], whose mean square is 2.5, so it becomes [0.632455,1.264911]. Head 0 of the second token is all zeros and stays zero because the 0.000001 inside the square root keeps the denominator positive.

Input: x = [[[0.5,-1,2,0,3,1]],[[1,1,1,-1,-1,-1]]], num_heads = 3

Output: {"heads":[[[[0.632455,-1.26491]],[[1.414213,0.0]],[[1.341641,0.447214]]],[[[1.0,1.0]],[[1.0,-1.0]],[[-1.0,-1.0]]]],"merged":[[[0.632455,-1.26491,1.414213,0.0,1.341641,0.447214]],[[1.0,1.0,1.0,-1.0,-1.0,-1.0]]]}

Input: x = [[[2,-2,1,0],[0.1,0.2,0.3,0.4],[0,0,0,0]]], num_heads = 1

Output: {"heads":[[[[1.333333,-1.333333,0.666667,0.0],[0.365146,0.730292,1.095438,1.460584],[0.0,0.0,0.0,0.0]]]],"merged":[[[1.333333,-1.333333,0.666667,0.0],[0.365146,0.730292,1.095438,1.460584],[0.0,0.0,0.0,0.0]]]}

Hint 1

jax.jit(f, static_argnames=("num_heads",)) passes num_heads as a plain Python int.

Hint 2

jnp.swapaxes(x, 1, 2) reorders axes, which a reshape alone never does.

Hint 3

jax.lax.rsqrt(v) is 1 / sqrt(v), and keepdims=True keeps a broadcastable axis.

Requirements

Compile the function with jax.jit with num_heads marked static, so that x is the only traced argument. The reshapes that split and merge heads need the head count as a Python integer while JAX traces the function.
The name per_head_rms_norm must refer to the compiled function itself, so that calling it by that name runs the jitted version.

Constraints

x is a float32 JAX array of shape (B, T, D) with B from 1 to 3, T from 1 to 4 and D from 2 to 12.
num_heads is a positive Python int that divides D.
All values of x are finite, and a head may be entirely zero.

Try Similar Problems

Traced Learning Rate ScheduleMedium Exponential Moving Average with ScanMedium Spectral Norm by Power IterationMedium Functional Array UpdateEasy PRNG Key SplittingEasy Masked Mean Without Boolean IndexingEasy Value and Grad of a LossEasy Model Ensemble with in_axesEasy SGD with Momentum over a PytreeEasy Rotary Position EmbeddingEasy KV Cache Update at a PositionEasy Device Mesh and Named ShardingEasy