| 253 | return t0, t1 |
| 254 | |
| 255 | def __get_sde_diffusion_and_drift( |
| 256 | self, |
| 257 | diffusion_form="SBDM", |
| 258 | diffusion_norm=1.0, |
| 259 | ): |
| 260 | |
| 261 | def diffusion_fn(x, t): |
| 262 | diffusion = self.schedule.compute_diffusion(x, t, form=diffusion_form, norm=diffusion_norm) |
| 263 | return diffusion |
| 264 | |
| 265 | sde_drift = \ |
| 266 | lambda x, t, model, **kwargs: \ |
| 267 | self.drift(x, t, model, **kwargs) + diffusion_fn(x, t) * self.score(x, t, model, **kwargs) |
| 268 | |
| 269 | sde_diffusion = diffusion_fn |
| 270 | |
| 271 | return sde_drift, sde_diffusion |
| 272 | |
| 273 | def last_step( |
| 274 | self, |