q_sample implements the 'nice property' — corrupting an image to any noise level in one operation

The q_sample function implements the closed-form forward diffusion: given a clean image x_start, a timestep batch t, and optional noise, it returns x_t = sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise. The precomputed tensors sqrt_alphas_cumprod and sqrt_one_minus_alphas_cumprod are indexed per batch element using an extract helper that gathers values at the right timestep indices. This single function replaces running the forward process t steps in sequence, and is the inner loop of training.

Examples

Code: def q_sample(x_start, t, noise=None): if noise is None: noise = torch.randn_like(x_start); return sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise.

Assessment

Given precomputed alpha_bar arrays, implement q_sample without looking at the tutorial. What does extract(a, t, x_shape) do and why is it needed?