Get the last step function of the SDE solver
(
self,
x,
t,
model,
sde_drift,
last_step,
last_step_size,
**model_kwargs
)
| 271 | return sde_drift, sde_diffusion |
| 272 | |
| 273 | def last_step( |
| 274 | self, |
| 275 | x, |
| 276 | t, |
| 277 | model, |
| 278 | sde_drift, |
| 279 | last_step, |
| 280 | last_step_size, |
| 281 | **model_kwargs |
| 282 | ): |
| 283 | """Get the last step function of the SDE solver""" |
| 284 | |
| 285 | if last_step is None: |
| 286 | return x |
| 287 | |
| 288 | elif last_step == "Mean": |
| 289 | return x + sde_drift(x, t, model, **model_kwargs) * last_step_size |
| 290 | |
| 291 | elif last_step == "Tweedie": |
| 292 | alpha = self.schedule.compute_alpha_t # simple aliasing; the original name was too long |
| 293 | sigma = self.schedule.compute_sigma_t |
| 294 | # return x / alpha(t)[0] + (sigma(t)[0] ** 2) / alpha(t)[0] * self.score(x, t, model, **model_kwargs) |
| 295 | raise NotImplementedError("Tweedie last step seems weird (alpha(t) is indexed twice?!?)") |
| 296 | |
| 297 | elif last_step == "Euler": |
| 298 | return x + self.drift(x, t, model, **model_kwargs) * last_step_size |
| 299 | |
| 300 | else: |
| 301 | raise NotImplementedError(f"Last step '{last_step}' not implemented.") |
| 302 | |
| 303 | def sample( |
| 304 | self, |