MolecularDiffusion.modules.models.syncogen.diffusion.sampling.discrete_strategies.utils¶
Functions¶
|
Sample from a categorical distribution with temperature scaling. |
|
Build a one‑hot edge tensor that keeps any already‑denoised incoming edge |
Module Contents¶
- MolecularDiffusion.modules.models.syncogen.diffusion.sampling.discrete_strategies.utils.sample_categorical(categorical_probs, temperature=1.0)¶
Sample from a categorical distribution with temperature scaling.
- Parameters:
categorical_probs – Unnormalized probabilities
temperature – Temperature parameter for scaling logits. Higher values produce more uniform samples.
- Returns:
Sampled indices from the categorical distribution normalized_probs: The normalized probabilities used for sampling
- Return type:
samples
- MolecularDiffusion.modules.models.syncogen.diffusion.sampling.discrete_strategies.utils.sample_edges(E: torch.Tensor, p_e0: torch.Tensor, lengths: torch.Tensor, argmax: bool = False) torch.Tensor¶
Build a one‑hot edge tensor that keeps any already‑denoised incoming edge and, for nodes still missing one, samples exactly ONE (parent, type) pair from the score logits. :param E: :type E: (B, n, n, R) current (partly‑denoised) one‑hot edges :param p_e0: :type p_e0: (B, n, n, R) score logits; last‑2 channels = [no‑edge, masked] :param lengths: :type lengths: (B,) number of real nodes per graph (padding after that) :param argmax: :type argmax: bool if True, take argmax instead of sampling (default False)