Study Plan

Exponential Moving Average with Scan

JAX

Medium

Compute the bias-corrected exponential moving average of a sequence, and return both the final running state and the corrected average at every step.

The rows of values are $x_1, \dots, x_T$ , each of shape (D,), and $\beta$ is decay. Starting from $m_0 = 0$ , for every step $t$ from 1 to T,

m_t = \beta \, m_{t-1} + (1 - \beta) \, x_t, \qquad \hat m_t = \frac{m_t}{1 - \beta^t}.

Here $m_t$ is the raw running average and $\hat m_t$ is its bias-corrected value.

Your function is compiled with jax.jit, with decay passed as a traced float32 scalar.

Return a dict with the key final_state, a float32 JAX array of shape (D,) holding the raw $m_T$ , and the key corrected, a float32 JAX array of shape (T, D) whose row $t - 1$ is $\hat m_t$ .

Examples

Input: values = [[1.0],[2.0],[3.0]], decay = 0.5

Output: {"final_state":[2.125],"corrected":[[1.0],[1.666667],[2.428571]]}

Explanation: The raw averages are 0.5, 1.25 and 2.125, and the weight sums 1 - 0.5^t are 0.5, 0.75 and 0.875, so the corrected averages are 1, 1.666667 and 2.428571. final_state is the raw value 2.125.

Input: values = [[2.0,-1.0],[2.0,-1.0],[2.0,-1.0],[2.0,-1.0]], decay = 0.9

Output: {"final_state":[0.6878,-0.3439],"corrected":[[2.0,-1.0],[2.0,-1.0],[2.0,-1.0],[2.0,-1.0]]}

Input: values = [[1.0,0.0,-1.0],[3.0,1.0,0.0]], decay = 0.0

Output: {"final_state":[3.0,1.0,0.0],"corrected":[[1.0,0.0,-1.0],[3.0,1.0,0.0]]}

Hint 1

final, outs = jax.lax.scan(step, init, xs) with step(carry, x) -> (new_carry, out).

Hint 2

The carry can be a tuple, such as (avg, power).

Hint 3

jnp.zeros_like(xs[0]) gives a starting state with the right shape and dtype.

Requirements

Write the loop over the T steps with jax.lax.scan and return the per-step results as its stacked outputs. A compiled program without such a scan, for example one built with a Python loop or jax.lax.fori_loop, is rejected.

Constraints

values is a float32 JAX array of shape (T, D) with T from 1 to 16 and D from 1 to 5.
decay is a float32 JAX scalar of shape () with a value from 0 to 0.99 inclusive.
All values are finite.

Try Similar Problems

Traced Learning Rate ScheduleMedium Spectral Norm by Power IterationMedium Static Arguments in jitEasy Pad and Mask to Bucket LengthMedium Hessian-Vector ProductMedium Per-Example Gradient ClippingMedium Straight-Through Estimator with Custom VJPMedium Batched Attention with vmapMedium Pairwise Distances with Nested vmapMedium Global Norm Gradient ClippingMedium AdamW from ScratchMedium Mixed Precision with Loss ScalingMedium