Multi-head attention as specified in https://arxiv.org/abs/2107.14795 Appendix E plus support for rotary position embeddings (https://arxiv.or

Pre-layer-norm cross-attention (see `MultiHeadAttention` for attention details).

Pre-layer norm self-attention (see `MultiHeadAttention` and for attention details).

Load the dataset

Load the dataset

Load the datasets

Load the datasets

Load the dataset

Load the dataset

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

input: grad_out: (m, c, nsample) output: (n, c), None

input: grad_out: (n, nsample, c) output: grad_input1: (n, c), grad_input2: (n, c)

input: grad_out: (n, c) output: grad_input: (n, c), grad_position: (n, nsample, c), grad_weight : (n, nsample, c')

input: xyz: (m, 3), new_xyz: (n, 3), input: (m, c), offset: (b), new_offset: (b) output: (n, c)

Evaluate a function without caching intermediate activations, allowing for reduced memory at the expense of extra compute in the backward pas

Close the board and flush the log

Close the board and flush the log

General collate function used for dataloader.

General collate function used for dataloader. This collate function is used for point-transformer

Create a 1D, 2D, or 3D convolution module.

Create a ScheduleSampler from a library of pre-defined samplers. :param name: the name of the sampler. :param diffusion: the diffusion o

Create visualizer according to the configuration Args: cfg: configuration dict Return: Visualizer

Sample x_{t+1} from the model using DDIM reverse ODE.

Generate samples from the model using DDIM. Same usage as p_sample_loop().

Denormalize pose sequence (usually for visualization) Args: pose_seq: a numpy array of normalized pose sequence

Denormalize the contact map Args: contact: the contact map to be denormalized clip: whether to clip the contact map

Denormalize the motion Args: motion: motion data to be denormalized Return: Denormalized motion

Denormalize pose sequence (usually for visualization) Args: pose_seq: a numpy array of normalized pose sequence

Denormalize the contact map Args: contact: the contact map to be denormalized clip: whether to clip the contact map

Denormalize the motion Args: motion: motion data to be denormalized Return: Denormalized motion

Evaluate sample Args: sample_list: sample result k_samples_list: k samples result save_dir: save directo

Evaluate sample Args: sample_list: sample result k_samples_list: k samples result save_dir: save directo

Evaluate sample Args: sample_list: sample result k_samples_list: k samples result save_dir: save directo

Evaluate sample Args: sample_list: sample result k_samples_list: k samples result save_dir: save directo

query (batch, query_dim) key (batch, seq_len, key_dim)