terminator.models.layers.utils.batchify

terminator.models.layers.utils.batchify(batched_flat_terms, term_lens)[source]

Take a flat representation of TERM information and batch them into a stacked representation.

In the TERM information condensor, TERM information is initially stored by concatenating all TERM tensors side by side in one dimension. However, for message passing, it’s convenient to batch these TERMs by splitting them and stacking them in a new dimension.

Parameters:

  • batched_flat_terms (torch.Tensor) – Tensor with shape (n_batch, sum_term_len, ...)

  • term_lens (list of (list of int)) – Length of each TERM per protein

Returns:

batchify_terms – Tensor with shape (n_batch, max_num_terms, max_term_len, ...)

Return type:

torch.Tensor