Gather consts for $t$ and reshape to feature map shape :param consts: (N, 1, 1) :param t: (N, H, W) :return: (N, H, W)
(consts: torch.Tensor, t: torch.Tensor)
| 37 | |
| 38 | |
| 39 | def gather(consts: torch.Tensor, t: torch.Tensor): |
| 40 | """ |
| 41 | Gather consts for $t$ and reshape to feature map shape |
| 42 | :param consts: (N, 1, 1) |
| 43 | :param t: (N, H, W) |
| 44 | :return: (N, H, W) |
| 45 | """ |
| 46 | c = consts.gather(-1, t) |
| 47 | return c.reshape(-1, 1, 1) |
| 48 | |
| 49 | |
| 50 | def q_xt_x0(x0, t, alpha_bar): |