A wrapper around torchsde.BrownianTree that enables batches of entropy.
| 63 | |
| 64 | |
| 65 | class BatchedBrownianTree: |
| 66 | """A wrapper around torchsde.BrownianTree that enables batches of entropy.""" |
| 67 | |
| 68 | def __init__(self, x, t0, t1, seed=None, **kwargs): |
| 69 | t0, t1, self.sign = self.sort(t0, t1) |
| 70 | w0 = kwargs.get('w0', torch.zeros_like(x)) |
| 71 | if seed is None: |
| 72 | seed = torch.randint(0, 2 ** 63 - 1, []).item() |
| 73 | self.batched = True |
| 74 | try: |
| 75 | assert len(seed) == x.shape[0] |
| 76 | w0 = w0[0] |
| 77 | except TypeError: |
| 78 | seed = [seed] |
| 79 | self.batched = False |
| 80 | self.trees = [torchsde.BrownianTree(t0, w0, t1, entropy=s, **kwargs) for s in seed] |
| 81 | |
| 82 | @staticmethod |
| 83 | def sort(a, b): |
| 84 | return (a, b, 1) if a < b else (b, a, -1) |
| 85 | |
| 86 | def __call__(self, t0, t1): |
| 87 | t0, t1, sign = self.sort(t0, t1) |
| 88 | w = torch.stack([tree(t0, t1) for tree in self.trees]) * (self.sign * sign) |
| 89 | return w if self.batched else w[0] |
| 90 | |
| 91 | |
| 92 | class BrownianTreeNoiseSampler: |