MolecularDiffusion.modules.models.tabasco.utils.tensor_ops¶
Functions¶
|
Zero out padded atoms. |
|
Center coordinates and zero padded atoms. |
|
Compute masked mean along the atom dimension. |
Module Contents¶
- MolecularDiffusion.modules.models.tabasco.utils.tensor_ops.apply_mask(x, padding_mask)¶
Zero out padded atoms.
- Parameters:
x – Tensor (*, N, D).
padding_mask – Bool/int tensor (*, N).
- Returns:
Masked tensor.
- Return type:
Tensor
Note: Based on the NVIDIA-Digital-Bio/proteina function of same name
- MolecularDiffusion.modules.models.tabasco.utils.tensor_ops.mask_and_zero_com(x, padding_mask)¶
Center coordinates and zero padded atoms.
- Parameters:
x – Tensor (*, N, D).
padding_mask – Bool/int tensor (*, N).
- Returns:
Centered and masked coordinates.
- Return type:
Tensor
Note: Based on the NVIDIA-Digital-Bio/proteina function of same name
- MolecularDiffusion.modules.models.tabasco.utils.tensor_ops.mean_w_mask(x, padding_mask)¶
Compute masked mean along the atom dimension.
Padded atoms (padding_mask == 1) are ignored. The result keeps a broadcastable shape (*, 1, D) and padded rows are zeroed for numerical safety.
- Parameters:
x – Tensor (*, N, D) of coordinates.
padding_mask – Bool/int tensor (*, N).
- Returns:
Mean coordinates with same dtype/device as x.
- Return type:
Tensor
Note: Based on the NVIDIA-Digital-Bio/proteina function of same name