mindsponge.cell.MSAColumnAttention

class mindsponge.cell.MSAColumnAttention(num_head, key_dim, gating, msa_act_dim, batch_size=None, slice_num=0)[source]

MSA column-wise gated self attention. The column-wise attention lets the elements that belong to the same target residue exchange information.

Reference:

Jumper et al. (2021) Suppl. Alg. 8 “MSAColumnAttention”.

Parameters

  • num_head (int) – The number of the heads.

  • key_dim (int) – The dimension of the input.

  • gating (bool) – Indicator of if the attention is gated.

  • msa_act_dim (int) – The dimension of the msa_act. The intermediate variable after MSA retrieving in AlphaFold.

  • batch_size (int) – The batch size of parameters in MSAColumnAttention, used in while control flow, Default: “None”.

  • slice_num (int) – The number of slices to be made to reduce memory, Default: 0.

Inputs:

  • msa_act (Tensor) - Tensor of msa_act. The intermediate variable after MSA retrieving in AlphaFold, shape \([N_{seqs}, N_{res}, C_m]\) .

  • msa_mask (Tensor) - The mask for MSAColumnAttention matrix, shape \([N_{seqs}, N_{res}]\).

  • index (Tensor) - The index of while loop, only used in case of while control flow. Default: “None”.

Outputs:

Tensor, the float tensor of the msa_act of the layer, shape \([N_{seqs}, N_{res}, C_m]\).

Supported Platforms:

Ascend GPU

Examples

>>> import numpy as np
>>> from mindsponge.cell import MSAColumnAttention
>>> from mindspore import dtype as mstype
>>> from mindspore import Tensor
>>> model = MSAColumnAttention(num_head=8, key_dim=256, gating=True,
...                         msa_act_dim=256, batch_size=1, slice_num=0)
>>> msa_act = Tensor(np.ones((512, 256, 256)), mstype.float32)
>>> msa_mask = Tensor(np.ones((512, 256)), mstype.float32)
>>> index = Tensor(0, mstype.int32)
>>> attn_out = model(msa_act, msa_mask, index)
>>> print(attn_out.shape)
(512, 256, 256)