Parameters ---------- x: torch.Tensor (batch, seqlen1, hidden_dim) (where hidden_dim = num heads * head dim)

Forward process. Args: x: tensor of shape [..., dim] Returns: embedding: an embedding of `x` of shape [...,

Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). This is the same as the DropConnect impl I created f

Args: pc (torch.FloatTensor): [B, N, 3] feats (torch.FloatTensor or None): [B, N, C] Returns:

Wrapper function: transfrom velocity prediction model to denoiser Args: velocity: [batch_dim, ...] shaped tensor; velocity model o

Wrapper function: transfrom score prediction model to velocity Args: score: [batch_dim, ...] shaped tensor; score model output

Return the argument as is.

Compute the Kullback-Leibler (KL) divergence between this distribution and another. If `other` is None, compute KL divergence to a s

Look up a symbol in a list of modules.

r""" Function that offloads all components, removes all model hooks that were added when using `enable_model_cpu_offload` and then app

Returns the real experiment on rank 0 and otherwise the DummyExperiment.

The pt_model is trained separately, so we already have access to its checkpoint and load it separately with `self.set_pt_model`.

process a single point cloud chunk, return the result and error info

Compute a distinct, deterministic RNG seed for each worker and node.

Repeatedly yield samples from an iterator.

Repeatedly returns batches from a DataLoader.

Resets parameters based on their initialization used in __init__.

Extract surface mesh using the Marching Cubes algorithm. Args: grid_logit (torch.Tensor): 3D grid logits tensor represen

Extract surface mesh using Differentiable Marching Cubes (DMC) algorithm. Args: grid_logit (torch.Tensor): 3D grid logit

Evaluate the given expression more safely.

returns a sampling function with given ODE settings and time step t Args: - sampling_method: type of sampler used in solving the ODE;

returns a sampling function for calculating likelihood with given ODE settings Args: - sampling_method: type of sampler used in solvin

returns a sampling function with given SDE settings Args: - sampling_method: type of sampler used in solving the SDE; default to be Eu

Forward process in flow-matching Args: sample (`torch.FloatTensor`): The input sample. times

Sets the begin index for the scheduler. This function should be run from pipeline before the inference. Args: begin_inde

Sets the begin index for the scheduler. This function should be run from pipeline before the inference. Args: begin_inde

Sets the discrete timesteps used for the diffusion chain (to be run before inference). Args: num_inference_steps (`int`)

Given a list of URLs, yields that list, possibly shuffled.

The index counter for current timestep. It will increase 1 after each scheduler step.

The index counter for current timestep. It will increase 1 after each scheduler step.

Given a stream of tar files, yield samples. Args: src: stream of tar files handler: exception handler select_files: funct