Appends dimensions to the end of a tensor until it has target_dims dimensions. From https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/utils.py
(x, target_dims)
| 3 | |
| 4 | |
| 5 | def append_dims(x, target_dims): |
| 6 | """Appends dimensions to the end of a tensor until it has target_dims dimensions. |
| 7 | From https://github.com/crowsonkb/k-diffusion/blob/master/k_diffusion/utils.py""" |
| 8 | dims_to_append = target_dims - x.ndim |
| 9 | if dims_to_append < 0: |
| 10 | raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less') |
| 11 | return x[(...,) + (None,) * dims_to_append] |
| 12 | |
| 13 | |
| 14 | def norm_thresholding(x0, value): |