(a, t, x_shape)
| 153 | # gaussian diffusion trainer class |
| 154 | |
| 155 | def extract(a, t, x_shape): |
| 156 | b, *_ = t.shape |
| 157 | out = a.gather(-1, t) |
| 158 | return out.reshape(b, *((1,) * (len(x_shape) - 1))) |
| 159 | |
| 160 | |
| 161 | def noise_like(shape, device, repeat=False): |
no outgoing calls
no test coverage detected