Compute the mean for the previous step, given a function cond_fn that computes the gradient of a conditional log probability with respect to x. In particular, cond_fn computes grad(log(p(y|x))), and we want to condition on y. This uses the conditioning strat
(self, cond_fn, p_mean_var, x, t, model_kwargs=None)
| 372 | ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) |
| 373 | |
| 374 | def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None): |
| 375 | """ |
| 376 | Compute the mean for the previous step, given a function cond_fn that |
| 377 | computes the gradient of a conditional log probability with respect to |
| 378 | x. In particular, cond_fn computes grad(log(p(y|x))), and we want to |
| 379 | condition on y. |
| 380 | |
| 381 | This uses the conditioning strategy from Sohl-Dickstein et al. (2015). |
| 382 | """ |
| 383 | gradient = cond_fn(x, t, **model_kwargs) |
| 384 | new_mean = p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float() |
| 385 | return new_mean |
| 386 | |
| 387 | def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None): |
| 388 | """ |