Closed-Form Forward Sample q(x_t | x_0)

Problem Statement

Jump straight from a clean sample $x_0$ to any noise level $t$ in a single operation, without simulating the intermediate steps.

Background

Applying the forward kernel $t$ times and repeatedly using the fact that the sum of two independent Gaussians is Gaussian collapses the whole chain into one closed form:

$q(x_t \mid x_0) = \mathcal{N}\!\left(x_t;\ \sqrt{\bar{\alpha}_t}\,x_0,\ (1-\bar{\alpha}_t) I\right)$

Via the reparameterization trick this becomes an explicit, differentiable formula:

$x_t = \sqrt{\bar{\alpha}_t}\,x_0 + \sqrt{1-\bar{\alpha}_t}\,\varepsilon, \qquad \varepsilon \sim \mathcal{N}(0, I)$

This is the single most-used line of code in diffusion training: it lets you sample a random timestep for every element of a batch and produce its noisy version in O(1), instead of running a t-step loop.

Note the coefficients are square roots of $\bar{\alpha}_t$ and $1-\bar{\alpha}_t$, and that they satisfy $(\sqrt{\bar{\alpha}_t})^2 + (\sqrt{1-\bar{\alpha}_t})^2 = 1$: the forward process is variance preserving.

Your Task

Implement:

def q_sample(x0, t, alpha_bars, noise):

Return the noisy sample $x_t$ as a NumPy array with the same shape as x0.

Input Format

• x0: NumPy array of any shape, the clean data.
• t (int): 0-based index into alpha_bars.
• alpha_bars: 1-D NumPy array of cumulative alphas.
• noise: NumPy array with the same shape as x0, the standard-normal sample $\varepsilon$.

Output Format

A NumPy array with the same shape as x0.

Sample

x0 = np.array([1.0, -1.0, 0.5])
noise = np.array([0.2, 0.4, -0.6])
ab = np.array([0.9, 0.5, 0.1])
print(np.round(q_sample(x0, 1, ab, noise), 4).tolist())

With $\bar{\alpha}_1 = 0.5$ both coefficients equal $\sqrt{0.5} \approx 0.7071$, so the output is 0.7071 * x0 + 0.7071 * noise.