Code-Paper Mismatch

[andylolu2]

Equation 6 & 7 from the paper suggests that the scores are computed from $\hat{x}_{t_i}$ (not $\hat{x}'_{t_i}$).

However, in the implementation, the update (Eq. 6) is applied to x then passed to update_fn, which will then compute the scores from the updated x ($\hat{x}'_{t_i}$ in the paper), instead of $\hat{x}_{t_i}$.

score_inverse_problems/cs.py

Lines 181 to 196 in b56e383

	def inpaint_update_fn(rng, state, x, t, mask, known, coeff):
	x_space = to_space(x)
	mean, std = sde.marginal_prob(known, t)
	rng, step_rng = jax.random.split(rng)
	noise = jax.random.normal(step_rng, x.shape)
	noise_space = to_space(noise)
	noisy_known = mean + batch_mul(std, noise_space)
	x_space = merge_known_with_mask(config, x_space, noisy_known, mask, coeff)
	x = from_space(x_space)
	rng, step_rng = jax.random.split(rng)
	x, x_mean = update_fn(step_rng, state, x, t)
	return x

Is there any reason to prefer one method over the other?