chore: import upstream snapshot with attribution

labml_nn/diffusion/ddpm/utils.py

@@ -0,0 +1,16 @@
+"""
+---
+title: Utility functions for DDPM experiment
+summary: >
+  Utility functions for DDPM experiment
+---
+
+# Utility functions for [DDPM](index.html) experiemnt
+"""
+import torch.utils.data
+
+
+def gather(consts: torch.Tensor, t: torch.Tensor):
+    """Gather consts for $t$ and reshape to feature map shape"""
+    c = consts.gather(-1, t)
+    return c.reshape(-1, 1, 1, 1)