| 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): |