| 103 | self.sampler_type = sampler_type |
| 104 | |
| 105 | def sample(self, x, model, **model_kwargs): |
| 106 | |
| 107 | device = x[0].device if isinstance(x, tuple) else x.device |
| 108 | |
| 109 | def _fn(t, x): |
| 110 | t = ( |
| 111 | th.ones(x[0].size(0)).to(device) * t |
| 112 | if isinstance(x, tuple) |
| 113 | else th.ones(x.size(0)).to(device) * t |
| 114 | ) |
| 115 | model_output = self.drift(x, t, model, **model_kwargs) |
| 116 | return model_output |
| 117 | |
| 118 | t = self.t.to(device) |
| 119 | |
| 120 | atol = [self.atol] * len(x) if isinstance(x, tuple) else [self.atol] |
| 121 | rtol = [self.rtol] * len(x) if isinstance(x, tuple) else [self.rtol] |
| 122 | samples = odeint(_fn, x, t, method=self.sampler_type, atol=atol, rtol=rtol) |
| 123 | return samples |