Restore the parameters stored with the `store` method. Useful to validate the model with EMA parameters without affecting the

Save the current parameters for restoring later. Args: parameters: Iterable of `torch.nn.Parameter`; the parameters to be

Create sinusoidal timestep embeddings. :param timesteps: a 1-D Tensor of N indices, one per batch element. These may be

Create a 1D, 2D, or 3D average pooling module.

imgs: [N, 1, num_voxels] imgs: [N, chan, T] pred: [N, L, p] mask: [N, L], 0 is keep, 1 is remove,

embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)

This matches the implementation in Denoising Diffusion Probabilistic Models: From Fairseq. Build sinusoidal embeddings. This matches

x, y: one dimension voxels array ratio: ratio for interpolation return: z same shape as x and y

Take the mean over all non-batch dimensions.

source: https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/losses.py#L12 Compute the K

imgs: (N, 1, num_voxels) imgs: [N, chan, T] x: (N, L, patch_size) x: [N, chan * 4, T/4]

Get the distribution q(x_t | x_0). :param x_start: the [N x C x ...] tensor of noiseless inputs. :param t: the number of diff

Perform per-sample random masking by per-sample shuffling. Per-sample shuffling is done by argsort random noise. x: [N, L, D]