Batch Utility: Jitted Versions#
Jitted versions of masking functions. Apparently, the traced versions are not faster than the original functions. This is likely due to the fact that the functions are already quite simple and do not benefit from the tracing optimizations.
- tad_mctc.batch.mask.jit.real_pairs_traced(numbers, mask_diagonal=True)[source]#
Create a mask for pairs of atoms from atomic numbers, discerning padding and actual atoms. Padding value is zero.
- Parameters:
numbers (Tensor) – Atomic numbers for all atoms.
mask_diagonal (bool, optional) – Flag for also masking the diagonal, i.e., all pairs with the same indices. Defaults to True, i.e., writing False to the diagonal.
- Returns:
Mask for atom pairs that discerns padding and real atoms.
- Return type:
Tensor