A noise sampler backed by a torchsde.BrownianTree. Args: x (Tensor): The tensor whose shape, device and dtype to use to generate random samples. sigma_min (float): The low end of the valid interval. sigma_max (float): The high end of the valid interval.
| 90 | |
| 91 | |
| 92 | class BrownianTreeNoiseSampler: |
| 93 | """A noise sampler backed by a torchsde.BrownianTree. |
| 94 | |
| 95 | Args: |
| 96 | x (Tensor): The tensor whose shape, device and dtype to use to generate |
| 97 | random samples. |
| 98 | sigma_min (float): The low end of the valid interval. |
| 99 | sigma_max (float): The high end of the valid interval. |
| 100 | seed (int or List[int]): The random seed. If a list of seeds is |
| 101 | supplied instead of a single integer, then the noise sampler will |
| 102 | use one BrownianTree per batch item, each with its own seed. |
| 103 | transform (callable): A function that maps sigma to the sampler's |
| 104 | internal timestep. |
| 105 | """ |
| 106 | |
| 107 | def __init__(self, x, sigma_min, sigma_max, seed=None, transform=lambda x: x): |
| 108 | self.transform = transform |
| 109 | t0, t1 = self.transform(torch.as_tensor(sigma_min)), self.transform(torch.as_tensor(sigma_max)) |
| 110 | self.tree = BatchedBrownianTree(x, t0, t1, seed) |
| 111 | |
| 112 | def __call__(self, sigma, sigma_next): |
| 113 | t0, t1 = self.transform(torch.as_tensor(sigma)), self.transform(torch.as_tensor(sigma_next)) |
| 114 | return self.tree(t0, t1) / (t1 - t0).abs().sqrt() |
| 115 | |
| 116 | |
| 117 | @torch.no_grad() |
no outgoing calls
no test coverage detected