MolecularDiffusion.modules.layers.functional

Functions

as_mask(indexes, length)

Convert indexes into a binary mask.

clipped_policy_gradient_objective(policy, agent, reward)

masked_mean(input, mask[, dim, keepdim])

Masked mean of a tensor.

mean_with_nan(input[, dim, keepdim])

Mean of a tensor. Ignore all nan values.

multi_slice(starts, ends)

Compute the union of indexes in multiple slices.

multi_slice_mask(starts, ends, length)

Compute the union of multiple slices into a binary mask.

multinomial(input, num_sample[, replacement])

Fast multinomial sampling. This is the default implementation in PyTorch v1.6.0+.

one_hot(index, size)

Expand indexes into one-hot vectors.

padded_to_variadic(padded, size)

Convert a padded tensor to a variadic tensor.

policy_gradient_objective(policy, reward)

shifted_softplus(input)

Shifted softplus function.

variadic_arange(size)

Return a 1-D tensor that contains integer intervals of variadic sizes.

variadic_cross_entropy(input, target, size[, reduction])

Compute cross entropy loss over categories with variadic sizes.

variadic_log_softmax(input, size)

Compute log softmax over categories with variadic sizes.

variadic_max(input, size)

Compute max over sets with variadic sizes.

variadic_mean(input, size)

Compute mean over sets with variadic sizes.

variadic_meshgrid(input1, size1, input2, size2)

Compute the Cartesian product for two batches of sets with variadic sizes.

variadic_randperm(size)

Return random permutations for sets with variadic sizes.

variadic_sample(input, size, num_sample)

Draw samples with replacement from sets with variadic sizes.

variadic_softmax(input, size)

Compute softmax over categories with variadic sizes.

variadic_sort(input, size[, descending])

Sort elements in sets with variadic sizes.

variadic_sum(input, size)

Compute sum over sets with variadic sizes.

variadic_to_padded(input, size[, value])

Convert a variadic tensor to a padded tensor.

variadic_topk(input, size, k[, largest])

Compute the \(k\) largest elements over sets with variadic sizes.

Module Contents

MolecularDiffusion.modules.layers.functional.as_mask(indexes, length)

Convert indexes into a binary mask.

Parameters:

  • indexes (LongTensor) – positive indexes

  • length (int) – maximal possible value of indexes

MolecularDiffusion.modules.layers.functional.clipped_policy_gradient_objective(policy, agent, reward, eps=0.2)
MolecularDiffusion.modules.layers.functional.masked_mean(input, mask, dim=None, keepdim=False)

Masked mean of a tensor.

Parameters:

  • input (Tensor) – input tensor

  • mask (BoolTensor) – mask tensor

  • dim (int or tuple of int, optional) – dimension to reduce

  • keepdim (bool, optional) – whether retain dim or not

MolecularDiffusion.modules.layers.functional.mean_with_nan(input, dim=None, keepdim=False)

Mean of a tensor. Ignore all nan values.

Parameters:

  • input (Tensor) – input tensor

  • dim (int or tuple of int, optional) – dimension to reduce

  • keepdim (bool, optional) – whether retain dim or not

MolecularDiffusion.modules.layers.functional.multi_slice(starts, ends)

Compute the union of indexes in multiple slices.

Example:

>>> mask = multi_slice(torch.tensor([0, 1, 4]), torch.tensor([2, 3, 6]), 6)
>>> assert (mask == torch.tensor([0, 1, 2, 4, 5]).all()
Parameters:

  • starts (LongTensor) – start indexes of slices

  • ends (LongTensor) – end indexes of slices

MolecularDiffusion.modules.layers.functional.multi_slice_mask(starts, ends, length)

Compute the union of multiple slices into a binary mask.

Example:

>>> mask = multi_slice_mask(torch.tensor([0, 1, 4]), torch.tensor([2, 3, 6]), 6)
>>> assert (mask == torch.tensor([1, 1, 1, 0, 1, 1])).all()
Parameters:

  • starts (LongTensor) – start indexes of slices

  • ends (LongTensor) – end indexes of slices

  • length (int) – length of mask

MolecularDiffusion.modules.layers.functional.multinomial(input, num_sample, replacement=False)

Fast multinomial sampling. This is the default implementation in PyTorch v1.6.0+.

Parameters:

  • input (Tensor) – unnormalized distribution

  • num_sample (int) – number of samples

  • replacement (bool, optional) – sample with replacement or not

MolecularDiffusion.modules.layers.functional.one_hot(index, size)

Expand indexes into one-hot vectors.

Parameters:

  • index (Tensor) – index

  • size (int) – size of the one-hot dimension

MolecularDiffusion.modules.layers.functional.padded_to_variadic(padded, size)

Convert a padded tensor to a variadic tensor.

Parameters:

  • padded (Tensor) – padded tensor of shape \((N, ...)\)

  • size (LongTensor) – size of sets of shape \((N,)\)

MolecularDiffusion.modules.layers.functional.policy_gradient_objective(policy, reward)
MolecularDiffusion.modules.layers.functional.shifted_softplus(input)

Shifted softplus function.

Parameters:

input (Tensor) – input tensor

MolecularDiffusion.modules.layers.functional.variadic_arange(size)

Return a 1-D tensor that contains integer intervals of variadic sizes. This is a variadic variant of torch.arange(stop).expand(batch_size, -1).

Suppose there are \(N\) intervals.

Parameters:

size (LongTensor) – size of intervals of shape \((N,)\)

MolecularDiffusion.modules.layers.functional.variadic_cross_entropy(input, target, size, reduction='mean')

Compute cross entropy loss over categories with variadic sizes.

Suppose there are \(N\) samples, and the numbers of categories in all samples are summed to \(B\).

Parameters:

  • input (Tensor) – prediction of shape \((B, ...)\)

  • target (Tensor) – target of shape \((N, ...)\). Each target is a relative index in a sample.

  • size (LongTensor) – number of categories of shape \((N,)\)

  • reduction (string, optional) – reduction to apply to the output. Available reductions are none, sum and mean.

MolecularDiffusion.modules.layers.functional.variadic_log_softmax(input, size)

Compute log softmax over categories with variadic sizes.

Suppose there are \(N\) samples, and the numbers of categories in all samples are summed to \(B\).

Parameters:

  • input (Tensor) – input of shape \((B, ...)\)

  • size (LongTensor) – number of categories of shape \((N,)\)

MolecularDiffusion.modules.layers.functional.variadic_max(input, size)

Compute max over sets with variadic sizes.

Suppose there are \(N\) sets, and the sizes of all sets are summed to \(B\).

Parameters:

  • input (Tensor) – input of shape \((B, ...)\)

  • size (LongTensor) – size of sets of shape \((N,)\)

Returns

(Tensor, LongTensor): max values and indexes

MolecularDiffusion.modules.layers.functional.variadic_mean(input, size)

Compute mean over sets with variadic sizes.

Suppose there are \(N\) sets, and the sizes of all sets are summed to \(B\).

Parameters:

  • input (Tensor) – input of shape \((B, ...)\)

  • size (LongTensor) – size of sets of shape \((N,)\)

MolecularDiffusion.modules.layers.functional.variadic_meshgrid(input1, size1, input2, size2)

Compute the Cartesian product for two batches of sets with variadic sizes.

Suppose there are \(N\) sets in each input, and the sizes of all sets are summed to \(B_1\) and \(B_2\) respectively.

Parameters:

  • input1 (Tensor) – input of shape \((B_1, ...)\)

  • size1 (LongTensor) – size of input1 of shape \((N,)\)

  • input2 (Tensor) – input of shape \((B_2, ...)\)

  • size2 (LongTensor) – size of input2 of shape \((N,)\)

Returns

(Tensor, Tensor): the first and the second elements in the Cartesian product

MolecularDiffusion.modules.layers.functional.variadic_randperm(size)

Return random permutations for sets with variadic sizes. The i-th permutation contains integers from 0 to size[i] - 1.

Suppose there are \(N\) sets.

Parameters:

  • size (LongTensor) – size of sets of shape \((N,)\)

  • device (torch.device, optional) – device of the tensor

MolecularDiffusion.modules.layers.functional.variadic_sample(input, size, num_sample)

Draw samples with replacement from sets with variadic sizes.

Suppose there are \(N\) sets, and the sizes of all sets are summed to \(B\).

Parameters:

  • input (Tensor) – input of shape \((B, ...)\)

  • size (LongTensor) – size of sets of shape \((N,)\)

  • num_sample (int) – number of samples to draw from each set

MolecularDiffusion.modules.layers.functional.variadic_softmax(input, size)

Compute softmax over categories with variadic sizes.

Suppose there are \(N\) samples, and the numbers of categories in all samples are summed to \(B\).

Parameters:

  • input (Tensor) – input of shape \((B, ...)\)

  • size (LongTensor) – number of categories of shape \((N,)\)

MolecularDiffusion.modules.layers.functional.variadic_sort(input, size, descending=False)

Sort elements in sets with variadic sizes.

Suppose there are \(N\) sets, and the sizes of all sets are summed to \(B\).

Parameters:

  • input (Tensor) – input of shape \((B, ...)\)

  • size (LongTensor) – size of sets of shape \((N,)\)

  • descending (bool, optional) – return ascending or descending order

Returns

(Tensor, LongTensor): sorted values and indexes

MolecularDiffusion.modules.layers.functional.variadic_sum(input, size)

Compute sum over sets with variadic sizes.

Suppose there are \(N\) sets, and the sizes of all sets are summed to \(B\).

Parameters:

  • input (Tensor) – input of shape \((B, ...)\)

  • size (LongTensor) – size of sets of shape \((N,)\)

MolecularDiffusion.modules.layers.functional.variadic_to_padded(input, size, value=0)

Convert a variadic tensor to a padded tensor.

Suppose there are \(N\) sets, and the sizes of all sets are summed to \(B\).

Parameters:

  • input (Tensor) – input of shape \((B, ...)\)

  • size (LongTensor) – size of sets of shape \((N,)\)

  • value (scalar) – fill value for padding

Returns:

padded tensor and mask

Return type:

(Tensor, BoolTensor)

MolecularDiffusion.modules.layers.functional.variadic_topk(input, size, k, largest=True)

Compute the \(k\) largest elements over sets with variadic sizes.

Suppose there are \(N\) sets, and the sizes of all sets are summed to \(B\).

If any set has less than than \(k\) elements, the size-th largest element will be repeated to pad the output to \(k\).

Parameters:

  • input (Tensor) – input of shape \((B, ...)\)

  • size (LongTensor) – size of sets of shape \((N,)\)

  • k (int or LongTensor) – the k in “top-k”. Can be a fixed value for all sets, or different values for different sets of shape \((N,)\).

  • largest (bool, optional) – return largest or smallest elements

Returns

(Tensor, LongTensor): top-k values and indexes