wangwei
/
community
Not watched
1
0
Fork 0
Code
Releases 0 Wiki Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain
You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.
Tree: b92421ea94
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'b92421ea94'
${ noResults }
community/reproduce/AlphaFold2-Chinese/module/basic_module.py

937 lines
53 kB
Raw Blame History 

  1. # Copyright 2021 Huawei Technologies Co., Ltd

  2. #

  3. # Licensed under the Apache License, Version 2.0 (the "License");

  4. # you may not use this file except in compliance with the License.

  5. # You may obtain a copy of the License at

  6. #

  7. # http://www.apache.org/licenses/LICENSE-2.0

  8. #

  9. # Unless required by applicable law or agreed to in writing, software

  10. # distributed under the License is distributed on an "AS IS" BASIS,

  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

  12. # See the License for the specific language governing permissions and

  13. # limitations under the License.

  14. # ============================================================================

  15. 

  16. """basic module"""

  17. 

  18. import mindspore.nn as nn

  19. import mindspore.common.dtype as mstype

  20. import mindspore.numpy as mnp

  21. from mindspore.ops import operations as P

  22. from mindspore.ops import functional as F

  23. from mindspore.common.tensor import Tensor

  24. from mindspore import Parameter

  25. import numpy as np

  26. from commons.utils import mask_mean

  27. 

  28. 

  29. class Attention(nn.Cell):

  30.     '''attention module'''

  31.     def __init__(self, config, q_data_dim, m_data_dim, output_dim, batch_size=None):

  32.         super(Attention, self).__init__()

  33.         self.config = config

  34.         self.q_data_dim = q_data_dim

  35.         self.m_data_dim = m_data_dim

  36.         self.output_dim = output_dim

  37.         self.num_head = self.config.num_head

  38.         self.gating = self.config.gating

  39.         self.key_dim = self.config.get('key_dim', int(q_data_dim))

  40.         self.value_dim = self.config.get('value_dim', int(m_data_dim))

  41.         self.key_dim = self.key_dim // self.num_head

  42.         self.value_dim = self.value_dim // self.num_head

  43.         self.batch_size = batch_size

  44.         self.matmul = P.MatMul(transpose_b=True)

  45.         self.batch_matmul_trans_b = P.BatchMatMul(transpose_b=True)

  46.         self.softmax = nn.Softmax()

  47.         self.sigmoid = nn.Sigmoid()

  48.         self.batch_size = batch_size

  49.         self._init_parameter()

  50. 

  51.     def _init_parameter(self):

  52.         '''init parameter'''

  53.         if self.batch_size:

  54.             self.linear_q_weights = Parameter(Tensor(np.zeros([self.batch_size, self.num_head * self.key_dim,

  55.                                                                self.q_data_dim]), mstype.float32))

  56.             self.linear_k_weights = Parameter(Tensor(np.zeros([self.batch_size, self.num_head * self.key_dim,

  57.                                                                self.m_data_dim]), mstype.float32))

  58.             self.linear_v_weights = Parameter(Tensor(np.zeros([self.batch_size, self.num_head * self.value_dim,

  59.                                                                self.m_data_dim]), mstype.float32))

  60.             self.linear_output_weights = Parameter(Tensor(np.zeros([self.batch_size, self.output_dim,

  61.                                                                     self.num_head * self.value_dim]), mstype.float32))

  62.             self.o_biases = Parameter(Tensor(np.zeros([self.batch_size, self.output_dim]), mstype.float32))

  63.             if self.gating:

  64.                 self.linear_gating_weights = Parameter(Tensor(np.zeros([self.batch_size, self.num_head * self.value_dim,

  65.                                                                         self.q_data_dim]), mstype.float32))

  66.                 self.gating_biases = Parameter(Tensor(np.zeros((self.batch_size, self.num_head, self.value_dim)),

  67.                                                       mstype.float32), name="gating_b")

  68.         else:

  69.             self.linear_q_weights = Parameter(Tensor(np.zeros([self.num_head * self.key_dim, self.q_data_dim]),

  70.                                                      mstype.float32))

  71.             self.linear_k_weights = Parameter(Tensor(np.zeros([self.num_head * self.key_dim, self.m_data_dim]),

  72.                                                      mstype.float32))

  73.             self.linear_v_weights = Parameter(Tensor(np.zeros([self.num_head * self.value_dim, self.m_data_dim]),

  74.                                                      mstype.float32))

  75.             self.linear_output_weights = Parameter(Tensor(np.zeros([self.output_dim, self.num_head * self.value_dim]),

  76.                                                           mstype.float32))

  77.             self.o_biases = Parameter(Tensor(np.zeros([self.output_dim]), mstype.float32))

  78.             if self.gating:

  79.                 self.linear_gating_weights = Parameter(Tensor(np.zeros([self.num_head * self.value_dim,

  80.                                                                         self.q_data_dim]), mstype.float32))

  81.                 self.gating_biases = Parameter(Tensor(np.zeros((self.num_head, self.value_dim)), mstype.float32),

  82.                                                name="gating_b")

  83. 

  84.     def construct(self, q_data, m_data, bias, index=None, nonbatched_bias=None):

  85.         '''construct'''

  86.         if self.batch_size:

  87.             linear_q_weight = P.Gather()(self.linear_q_weights, index, 0)

  88.             linear_k_weight = P.Gather()(self.linear_k_weights, index, 0)

  89.             linear_v_weight = P.Gather()(self.linear_v_weights, index, 0)

  90.             linear_output_weight = P.Gather()(self.linear_output_weights, index, 0)

  91.             o_bias = P.Gather()(self.o_biases, index, 0)

  92.             linear_gating_weight = 0

  93.             gating_bias = 0

  94.             if self.gating:

  95.                 linear_gating_weight = P.Gather()(self.linear_gating_weights, index, 0)

  96.                 gating_bias = P.Gather()(self.gating_biases, index, 0)

  97.         else:

  98.             linear_q_weight = self.linear_q_weights

  99.             linear_k_weight = self.linear_k_weights

  100.             linear_v_weight = self.linear_v_weights

  101.             linear_output_weight = self.linear_output_weights

  102.             o_bias = self.o_biases

  103.             linear_gating_weight = 0

  104.             gating_bias = 0

  105.             if self.gating:

  106.                 linear_gating_weight = self.linear_gating_weights

  107.                 gating_bias = self.gating_biases

  108. 

  109.         q_data, m_data, bias = q_data.astype(mstype.float32), m_data.astype(mstype.float32), bias.astype(mstype.float32)

  110.         dim_b, dim_q, dim_a = q_data.shape

  111.         _, dim_k, dim_c = m_data.shape

  112.         dim_h = self.num_head

  113. 

  114.         q_data = P.Reshape()(q_data, (-1, dim_a))

  115.         m_data = P.Reshape()(m_data, (-1, dim_c))

  116. 

  117.         q = self.matmul(q_data.astype(mstype.float16), linear_q_weight.astype(mstype.float16)). \

  118.                 astype(mstype.float32) * self.key_dim ** (-0.5)

  119.         k = self.matmul(m_data.astype(mstype.float16), linear_k_weight.astype(mstype.float16))

  120.         v = self.matmul(m_data.astype(mstype.float16), linear_v_weight.astype(mstype.float16))

  121. 

  122.         q = P.Reshape()(q, (dim_b, dim_q, dim_h, -1))

  123.         k = P.Reshape()(k, (dim_b, dim_k, dim_h, -1))

  124.         v = P.Reshape()(v, (dim_b, dim_k, dim_h, -1))

  125. 

  126.         tmp_q = P.Reshape()(P.Transpose()(q.astype(mstype.float16), (0, 2, 1, 3)), (dim_b * dim_h, dim_q, -1))

  127.         tmp_k = P.Reshape()(P.Transpose()(k.astype(mstype.float16), (0, 2, 1, 3)), (dim_b * dim_h, dim_k, -1))

  128.         logits = P.Reshape()(self.batch_matmul_trans_b(tmp_q.astype(mstype.float16),

  129.                                                        tmp_k.astype(mstype.float16)),

  130.                              (dim_b, dim_h, dim_q, dim_k)) + bias.astype(mstype.float16)

  131. 

  132.         if nonbatched_bias is not None:

  133.             logits += mnp.expand_dims(nonbatched_bias, axis=0)

  134.         weights = self.softmax(logits.astype(mstype.float32))

  135.         tmp_v = P.Reshape()(P.Transpose()(v, (0, 2, 3, 1)), (dim_b * dim_h, -1, dim_k))

  136.         tmp_weights = P.Reshape()(weights, (dim_b * dim_h, dim_q, -1))

  137.         weighted_avg = P.Transpose()(P.Reshape()(self.batch_matmul_trans_b(tmp_weights.astype(mstype.float16),

  138.                                                                            tmp_v.astype(mstype.float16)),

  139.                                                  (dim_b, dim_h, dim_q, -1)),

  140.                                      (0, 2, 1, 3)).astype(mstype.float32)

  141. 

  142.         if self.gating:

  143.             gate_values = P.Reshape()(

  144.                 self.matmul(q_data.astype(mstype.float16), linear_gating_weight.astype(mstype.float16)),

  145.                 (dim_b, dim_q, dim_h, -1)) + gating_bias[None, None, ...].astype(mstype.float16)

  146.             gate_values = self.sigmoid(gate_values.astype(mstype.float32))

  147.             weighted_avg = P.Reshape()(weighted_avg * gate_values, (dim_b * dim_q, -1))

  148. 

  149.         weighted_avg = P.Reshape()(weighted_avg, (dim_b * dim_q, -1))

  150.         output = P.Reshape()(

  151.             self.matmul(weighted_avg.astype(mstype.float16), linear_output_weight.astype(mstype.float16)),

  152.             (dim_b, dim_q, -1)) + o_bias[None, ...].astype(mstype.float16)

  153.         return output

  154. 

  155. 

  156. class MSARowAttentionWithPairBias(nn.Cell):

  157.     '''MSA row attention'''

  158.     def __init__(self, config, msa_act_dim, pair_act_dim, batch_size=None, slice_num=0):

  159.         super(MSARowAttentionWithPairBias, self).__init__()

  160.         self.config = config

  161.         self.num_head = self.config.num_head

  162.         self.batch_size = batch_size

  163.         self.norm = P.LayerNorm(begin_norm_axis=-1, begin_params_axis=-1, epsilon=1e-5)

  164.         self.matmul = P.MatMul(transpose_b=True)

  165.         self.attn_mod = Attention(self.config, msa_act_dim, msa_act_dim, msa_act_dim, batch_size)

  166.         self.msa_act_dim = msa_act_dim

  167.         self.pair_act_dim = pair_act_dim

  168.         self.batch_size = batch_size

  169.         self.slice_num = slice_num

  170.         self.idx = Tensor(0, mstype.int32)

  171.         self._init_parameter()

  172. 

  173.     def _init_parameter(self):

  174.         '''init parameter'''

  175.         self.query_norm_gammas = Parameter(Tensor(np.zeros([self.batch_size, self.msa_act_dim,]), mstype.float32))

  176.         self.query_norm_betas = Parameter(Tensor(np.zeros([self.batch_size, self.msa_act_dim,]), mstype.float32))

  177.         self.feat_2d_norm_gammas = Parameter(Tensor(np.zeros([self.batch_size, self.pair_act_dim,]), mstype.float32))

  178.         self.feat_2d_norm_betas = Parameter(Tensor(np.zeros([self.batch_size, self.pair_act_dim,]), mstype.float32))

  179.         self.feat_2d_weights = Parameter(

  180.             Tensor(np.zeros([self.batch_size, self.num_head, self.pair_act_dim]), mstype.float32))

  181. 

  182.     def construct(self, msa_act, msa_mask, pair_act, index):

  183.         '''construct'''

  184.         query_norm_gamma = P.Gather()(self.query_norm_gammas, index, 0)

  185.         query_norm_beta = P.Gather()(self.query_norm_betas, index, 0)

  186.         feat_2d_norm_gamma = P.Gather()(self.feat_2d_norm_gammas, index, 0)

  187.         feat_2d_norm_beta = P.Gather()(self.feat_2d_norm_betas, index, 0)

  188.         feat_2d_weight = P.Gather()(self.feat_2d_weights, index, 0)

  189. 

  190.         q, k, _ = pair_act.shape

  191.         bias = (1e9 * (msa_mask - 1.))[:, None, None, :]

  192.         msa_act, _, _ = self.norm(msa_act.astype(mstype.float32), query_norm_gamma.astype(mstype.float32),

  193.                                   query_norm_beta.astype(mstype.float32))

  194.         pair_act, _, _ = self.norm(pair_act.astype(mstype.float32), feat_2d_norm_gamma.astype(mstype.float32),

  195.                                    feat_2d_norm_beta.astype(mstype.float32))

  196.         pair_act = P.Reshape()(pair_act, (-1, pair_act.shape[-1]))

  197.         nonbatched_bias = P.Transpose()(

  198.             P.Reshape()(self.matmul(pair_act.astype(mstype.float16), feat_2d_weight.astype(mstype.float16)),

  199.                         (q, k, self.num_head)), (2, 0, 1))

  200. 

  201.         if self.slice_num:

  202.             msa_act_ori_shape = P.Shape()(msa_act)

  203.             slice_shape = (self.slice_num, -1) + msa_act_ori_shape[1:]

  204.             msa_act = P.Reshape()(msa_act, slice_shape).astype(mstype.float16)

  205.             bias_shape = P.Shape()(bias)

  206.             bias = P.Reshape()(bias, slice_shape[:2] + bias_shape[1:])

  207.             slice_idx = 0

  208.             slice_idx_tensor = self.idx

  209.             msa_act_tuple = ()

  210. 

  211.             msa_act_slice = P.Gather()(msa_act, slice_idx_tensor, 0)

  212.             bias_slice = P.Gather()(bias, slice_idx_tensor, 0)

  213.             msa_act_slice = self.attn_mod(msa_act_slice, msa_act_slice, bias_slice, index, nonbatched_bias)

  214.             msa_act_slice = P.Reshape()(msa_act_slice, ((1,) + P.Shape()(msa_act_slice)))

  215.             msa_act_tuple = msa_act_tuple + (msa_act_slice,)

  216.             slice_idx += 1

  217.             slice_idx_tensor += 1

  218. 

  219.             while slice_idx < self.slice_num:

  220.                 msa_act_slice = P.Gather()(msa_act, slice_idx_tensor, 0)

  221.                 msa_act_slice = F.depend(msa_act_slice, msa_act_tuple[-1])

  222.                 bias_slice = P.Gather()(bias, slice_idx_tensor, 0)

  223.                 msa_act_slice = self.attn_mod(msa_act_slice, msa_act_slice, bias_slice, index, nonbatched_bias)

  224.                 msa_act_slice = P.Reshape()(msa_act_slice, ((1,) + P.Shape()(msa_act_slice)))

  225.                 msa_act_tuple = msa_act_tuple + (msa_act_slice,)

  226.                 slice_idx += 1

  227.                 slice_idx_tensor += 1

  228. 

  229.             msa_act = P.Concat()(msa_act_tuple)

  230.             msa_act = P.Reshape()(msa_act, msa_act_ori_shape)

  231.             return msa_act

  232. 

  233.         msa_act = self.attn_mod(msa_act, msa_act, bias, index, nonbatched_bias)

  234.         return msa_act

  235. 

  236. 

  237. class MSAColumnAttention(nn.Cell):

  238.     '''MSA column attention'''

  239.     def __init__(self, config, msa_act_dim, batch_size=None, slice_num=0):

  240.         super(MSAColumnAttention, self).__init__()

  241.         self.config = config

  242.         self.query_norm = P.LayerNorm(begin_norm_axis=-1, begin_params_axis=-1, epsilon=1e-5)

  243.         self.attn_mod = Attention(self.config, msa_act_dim, msa_act_dim, msa_act_dim, batch_size)

  244.         self.batch_size = batch_size

  245.         self.slice_num = slice_num

  246.         self.msa_act_dim = msa_act_dim

  247.         self.idx = Tensor(0, mstype.int32)

  248.         self._init_parameter()

  249. 

  250.     def _init_parameter(self):

  251.         self.query_norm_gammas = Parameter(Tensor(np.zeros([self.batch_size, self.msa_act_dim]), mstype.float32))

  252.         self.query_norm_betas = Parameter(Tensor(np.zeros([self.batch_size, self.msa_act_dim]), mstype.float32))

  253. 

  254.     def construct(self, msa_act, msa_mask, index):

  255.         '''construct'''

  256.         query_norm_gamma = P.Gather()(self.query_norm_gammas, index, 0)

  257.         query_norm_beta = P.Gather()(self.query_norm_betas, index, 0)

  258.         msa_act = mnp.swapaxes(msa_act, -2, -3)

  259.         msa_mask = mnp.swapaxes(msa_mask, -1, -2)

  260.         bias = (1e9 * (msa_mask - 1.))[:, None, None, :]

  261.         msa_act, _, _ = self.query_norm(msa_act.astype(mstype.float32), query_norm_gamma.astype(mstype.float32),

  262.                                         query_norm_beta.astype(mstype.float32))

  263.         if self.slice_num:

  264.             msa_act_ori_shape = P.Shape()(msa_act)

  265.             slice_shape = (self.slice_num, -1) + msa_act_ori_shape[1:]

  266.             msa_act = P.Reshape()(msa_act, slice_shape).astype(mstype.float16)

  267.             bias_shape = P.Shape()(bias)

  268.             bias = P.Reshape()(bias, slice_shape[:2] + bias_shape[1:])

  269. 

  270.             slice_idx = 0

  271.             slice_idx_tensor = self.idx

  272.             msa_act_tuple = ()

  273. 

  274.             msa_act_slice = P.Gather()(msa_act, slice_idx_tensor, 0)

  275.             bias_slice = P.Gather()(bias, slice_idx_tensor, 0)

  276.             msa_act_slice = self.attn_mod(msa_act_slice, msa_act_slice, bias_slice, index)

  277.             msa_act_slice = P.Reshape()(msa_act_slice, ((1,) + P.Shape()(msa_act_slice)))

  278.             msa_act_tuple = msa_act_tuple + (msa_act_slice,)

  279.             slice_idx += 1

  280.             slice_idx_tensor += 1

  281. 

  282.             while slice_idx < self.slice_num:

  283.                 msa_act_slice = P.Gather()(msa_act, slice_idx_tensor, 0)

  284.                 msa_act_slice = F.depend(msa_act_slice, msa_act_tuple[-1])

  285.                 bias_slice = P.Gather()(bias, slice_idx_tensor, 0)

  286.                 msa_act_slice = self.attn_mod(msa_act_slice, msa_act_slice, bias_slice, index)

  287.                 msa_act_slice = P.Reshape()(msa_act_slice, ((1,) + P.Shape()(msa_act_slice)))

  288.                 msa_act_tuple = msa_act_tuple + (msa_act_slice,)

  289.                 slice_idx += 1

  290.                 slice_idx_tensor += 1

  291. 

  292.             msa_act = P.Concat()(msa_act_tuple)

  293.             msa_act = P.Reshape()(msa_act, msa_act_ori_shape)

  294.             msa_act = mnp.swapaxes(msa_act, -2, -3)

  295.             return msa_act

  296. 

  297.         msa_act = self.attn_mod(msa_act, msa_act, bias, index)

  298.         msa_act = mnp.swapaxes(msa_act, -2, -3)

  299.         return msa_act

  300. 

  301. 

  302. class GlobalAttention(nn.Cell):

  303.     '''global attention'''

  304.     def __init__(self, config, key_dim, value_dim, output_dim, batch_size=None):

  305.         super(GlobalAttention, self).__init__()

  306.         self.config = config

  307.         self.key_dim = key_dim

  308.         self.ori_key_dim = key_dim

  309.         self.value_dim = value_dim

  310.         self.ori_value_dim = value_dim

  311.         self.num_head = self.config.num_head

  312.         self.key_dim = self.key_dim // self.num_head

  313.         self.value_dim = self.value_dim // self.num_head

  314.         self.output_dim = output_dim

  315.         self.matmul_trans_b = P.MatMul(transpose_b=True)

  316.         self.batch_matmul = P.BatchMatMul()

  317.         self.batch_matmul_trans_b = P.BatchMatMul(transpose_b=True)

  318.         self.matmul = P.MatMul()

  319.         self.softmax = nn.Softmax()

  320.         self.sigmoid = nn.Sigmoid()

  321.         self.gating = self.config.gating

  322. 

  323.         self.batch_size = batch_size

  324.         self._init_parameter()

  325. 

  326.     def _init_parameter(self):

  327.         '''init parameter'''

  328.         self.linear_q_weights = Parameter(

  329.             Tensor(np.zeros((self.batch_size, self.ori_key_dim, self.num_head, self.key_dim)), mstype.float32))

  330.         self.linear_k_weights = Parameter(

  331.             Tensor(np.zeros((self.batch_size, self.ori_value_dim, self.key_dim)), mstype.float32))

  332.         self.linear_v_weights = Parameter(

  333.             Tensor(np.zeros((self.batch_size, self.ori_value_dim, self.value_dim)), mstype.float32))

  334.         self.linear_output_weights = Parameter(

  335.             Tensor(np.zeros((self.batch_size, self.output_dim, self.num_head * self.value_dim)), mstype.float32))

  336.         self.o_biases = Parameter(Tensor(np.zeros((self.batch_size, self.output_dim)), mstype.float32))

  337.         if self.gating:

  338.             self.linear_gating_weights = Parameter(

  339.                 Tensor(np.zeros((self.batch_size, self.num_head * self.value_dim, self.ori_key_dim)), mstype.float32))

  340.             self.gating_biases = Parameter(Tensor(np.zeros((self.batch_size, self.ori_key_dim)), mstype.float32))

  341. 

  342.     def construct(self, q_data, m_data, q_mask, bias, index):

  343.         '''construct'''

  344.         q_weights = P.Gather()(self.linear_q_weights, index, 0)

  345.         k_weights = P.Gather()(self.linear_k_weights, index, 0)

  346.         v_weights = P.Gather()(self.linear_v_weights, index, 0)

  347.         output_weights = P.Gather()(self.linear_output_weights, index, 0)

  348.         output_bias = P.Gather()(self.o_biases, index, 0)

  349.         gating_weights = 0

  350.         gating_bias = 0

  351.         if self.gating:

  352.             gating_weights = P.Gather()(self.linear_gating_weights, index, 0)

  353.             gating_bias = P.Gather()(self.gating_biases, index, 0)

  354.         b, _, _ = m_data.shape

  355.         v_weights = v_weights[None, ...]

  356.         v_weights = mnp.broadcast_to(v_weights, (b, self.value_dim * self.num_head, self.value_dim))

  357.         v = self.batch_matmul(m_data.astype(mstype.float16), v_weights.astype(mstype.float16))

  358.         q_avg = mask_mean(q_mask, q_data, axis=1)

  359.         q_weights = P.Reshape()(q_weights, (-1, self.num_head * self.key_dim))

  360.         q = P.Reshape()(self.matmul(q_avg.astype(mstype.float16), q_weights.astype(mstype.float16)),

  361.                         (-1, self.num_head, self.key_dim)) * (self.key_dim ** (-0.5))

  362. 

  363.         k_weights = k_weights[None, ...]

  364.         k_weights = mnp.broadcast_to(k_weights, (b, self.value_dim * self.num_head, self.key_dim))

  365. 

  366.         k = self.batch_matmul(m_data.astype(mstype.float16), k_weights.astype(mstype.float16))

  367. 

  368.         bias = (1e9 * (q_mask[:, None, :, 0] - 1.))

  369. 

  370.         logits = self.batch_matmul_trans_b(q.astype(mstype.float16), k.astype(mstype.float16)) + bias.astype(

  371.             mstype.float16)

  372.         weights = self.softmax(logits.astype(mstype.float32))

  373.         weighted_avg = self.batch_matmul(weights.astype(mstype.float16), v.astype(mstype.float16))

  374. 

  375.         if self.gating:

  376.             # gate_values = self.linear_gating(q_data).astype(mstype.float32)

  377.             q_data_shape = P.Shape()(q_data)

  378.             if len(q_data_shape) != 2:

  379.                 q_data = P.Reshape()(q_data, (-1, q_data_shape[-1]))

  380.             out_shape = q_data_shape[:-1] + (-1,)

  381.             gate_values = P.Reshape()(self.matmul_trans_b(q_data.astype(mstype.float16),

  382.                                                           gating_weights.astype(mstype.float16)) +

  383.                                       gating_bias.astype(mstype.float16), out_shape)

  384. 

  385.             gate_values = P.Reshape()(self.sigmoid(gate_values.astype(mstype.float32)),

  386.                                       (b, -1, self.num_head, self.value_dim))

  387.             weighted_avg = P.Reshape()(weighted_avg[:, None] * gate_values, (-1, self.num_head * self.value_dim))

  388.             weighted_avg_shape = P.Shape()(weighted_avg)

  389.             if len(weighted_avg_shape) != 2:

  390.                 weighted_avg = P.Reshape()(weighted_avg, (-1, weighted_avg_shape[-1]))

  391.             output = P.Reshape()(self.matmul_trans_b(weighted_avg.astype(mstype.float16),

  392.                                                      output_weights.astype(mstype.float16))

  393.                                  + output_bias.astype(mstype.float16), (b, -1, self.output_dim))

  394. 

  395.         else:

  396.             weighted_avg = P.Reshape()(weighted_avg, (-1, self.num_head * self.value_dim))

  397.             # output = self.linear_gating(weighted_avg)

  398.             weighted_avg_shape = P.Shape()(weighted_avg)

  399.             if len(weighted_avg_shape) != 2:

  400.                 weighted_avg = P.Reshape()(weighted_avg, (-1, weighted_avg_shape[-1]))

  401.             out_shape = weighted_avg_shape[:-1] + (-1,)

  402.             output = P.Reshape()(self.matmul_trans_b(weighted_avg.astype(mstype.float16),

  403.                                                      output_weights.astype(mstype.float16)) +

  404.                                  output_bias.astype(mstype.float16), out_shape)

  405.             output = output[:, None]

  406.         return output

  407. 

  408. 

  409. class MSAColumnGlobalAttention(nn.Cell):

  410.     '''MSA column global attention'''

  411.     def __init__(self, config, msa_act_dim, batch_size=None, slice_num=0):

  412.         super(MSAColumnGlobalAttention, self).__init__()

  413.         self.config = config

  414.         self.attn_mod = GlobalAttention(self.config, msa_act_dim, msa_act_dim, msa_act_dim, batch_size)

  415.         self.query_norm = P.LayerNorm(begin_norm_axis=-1, begin_params_axis=-1, epsilon=1e-5)

  416.         self.batch_size = batch_size

  417.         self.slice_num = slice_num

  418.         self.msa_act_dim = msa_act_dim

  419.         self.idx = Tensor(0, mstype.int32)

  420.         self._init_parameter()

  421. 

  422.     def _init_parameter(self):

  423.         '''init parameter'''

  424.         self.query_norm_gammas = Parameter(Tensor(np.zeros((self.batch_size, self.msa_act_dim)), mstype.float32))

  425.         self.query_norm_betas = Parameter(Tensor(np.zeros((self.batch_size, self.msa_act_dim)), mstype.float32))

  426. 

  427.     def construct(self, msa_act, msa_mask, index):

  428.         '''construct'''

  429.         query_norm_gamma = P.Gather()(self.query_norm_gammas, index, 0)

  430.         query_norm_beta = P.Gather()(self.query_norm_betas, index, 0)

  431.         msa_act = mnp.swapaxes(msa_act, -2, -3)

  432.         msa_mask = mnp.swapaxes(msa_mask, -1, -2)

  433.         bias = (1e9 * (msa_mask - 1.))[:, None, None, :]

  434.         msa_act, _, _ = self.query_norm(msa_act.astype(mstype.float32),

  435.                                         query_norm_gamma.astype(mstype.float32),

  436.                                         query_norm_beta.astype(mstype.float32))

  437.         msa_mask = mnp.expand_dims(msa_mask, axis=-1)

  438. 

  439.         if self.slice_num:

  440.             msa_act_ori_shape = P.Shape()(msa_act)

  441.             slice_shape = (self.slice_num, -1) + msa_act_ori_shape[1:]

  442.             msa_act = P.Reshape()(msa_act, slice_shape).astype(mstype.float16)

  443.             bias_shape = P.Shape()(bias)

  444.             bias = P.Reshape()(bias, slice_shape[:2] + bias_shape[1:])

  445.             msa_mask_shape = P.Shape()(msa_mask)

  446.             msa_mask = P.Reshape()(msa_mask, slice_shape[:2] + msa_mask_shape[1:])

  447. 

  448.             slice_idx = 0

  449.             slice_idx_tensor = self.idx

  450.             msa_act_tuple = ()

  451. 

  452.             msa_act_slice = P.Gather()(msa_act, slice_idx_tensor, 0)

  453.             msa_mask_slice = P.Gather()(msa_mask, slice_idx_tensor, 0)

  454.             bias_slice = P.Gather()(bias, slice_idx_tensor, 0)

  455.             msa_act_slice = self.attn_mod(msa_act_slice, msa_act_slice, msa_mask_slice, bias_slice, index)

  456.             msa_act_slice = P.Reshape()(msa_act_slice, ((1,) + P.Shape()(msa_act_slice)))

  457.             msa_act_tuple = msa_act_tuple + (msa_act_slice,)

  458.             slice_idx += 1

  459.             slice_idx_tensor += 1

  460. 

  461.             while slice_idx < self.slice_num:

  462.                 msa_act_slice = P.Gather()(msa_act, slice_idx_tensor, 0)

  463.                 msa_act_slice = F.depend(msa_act_slice, msa_act_tuple[-1])

  464.                 msa_mask_slice = P.Gather()(msa_mask, slice_idx_tensor, 0)

  465.                 bias_slice = P.Gather()(bias, slice_idx_tensor, 0)

  466. 

  467.                 msa_act_slice = self.attn_mod(msa_act_slice, msa_act_slice, msa_mask_slice, bias_slice, index)

  468.                 msa_act_slice = P.Reshape()(msa_act_slice, ((1,) + P.Shape()(msa_act_slice)))

  469.                 msa_act_tuple = msa_act_tuple + (msa_act_slice,)

  470.                 slice_idx += 1

  471.                 slice_idx_tensor += 1

  472. 

  473.             msa_act = P.Concat()(msa_act_tuple)

  474.             msa_act = P.Reshape()(msa_act, msa_act_ori_shape)

  475.             msa_act = mnp.swapaxes(msa_act, -2, -3)

  476.             return msa_act

  477. 

  478.         msa_act = self.attn_mod(msa_act, msa_act, msa_mask, bias, index)

  479.         msa_act = mnp.swapaxes(msa_act, -2, -3)

  480.         return msa_act

  481. 

  482. 

  483. class Transition(nn.Cell):

  484.     '''transition'''

  485.     def __init__(self, config, layer_norm_dim, batch_size=None, slice_num=0):

  486.         super(Transition, self).__init__()

  487.         self.config = config

  488.         self.input_layer_norm = P.LayerNorm(begin_norm_axis=-1, begin_params_axis=-1, epsilon=1e-5)

  489.         self.matmul = P.MatMul(transpose_b=True)

  490.         self.layer_norm_dim = layer_norm_dim

  491.         self.num_intermediate = int(layer_norm_dim * self.config.num_intermediate_factor)

  492.         self.batch_size = batch_size

  493.         self.slice_num = slice_num

  494.         self.relu = nn.ReLU()

  495.         self.idx = Tensor(0, mstype.int32)

  496.         self._init_parameter()

  497. 

  498.     def _init_parameter(self):

  499.         '''init parameter'''

  500.         self.input_layer_norm_gammas = Parameter(

  501.             Tensor(np.zeros((self.batch_size, self.layer_norm_dim)), mstype.float32))

  502.         self.input_layer_norm_betas = Parameter(

  503.             Tensor(np.zeros((self.batch_size, self.layer_norm_dim)), mstype.float32))

  504.         self.transition1_weights = Parameter(

  505.             Tensor(np.zeros((self.batch_size, self.num_intermediate, self.layer_norm_dim)), mstype.float32))

  506.         self.transition1_biases = Parameter(Tensor(np.zeros((self.batch_size, self.num_intermediate)), mstype.float32))

  507.         self.transition2_weights = Parameter(

  508.             Tensor(np.zeros((self.batch_size, self.layer_norm_dim, self.num_intermediate)), mstype.float32))

  509.         self.transition2_biases = Parameter(Tensor(np.zeros((self.batch_size, self.layer_norm_dim)), mstype.float32))

  510. 

  511.     def construct(self, act, index):

  512.         '''construct'''

  513.         input_layer_norm_gamma = P.Gather()(self.input_layer_norm_gammas, index, 0)

  514.         input_layer_norm_beta = P.Gather()(self.input_layer_norm_betas, index, 0)

  515.         transition1_weight = P.Gather()(self.transition1_weights, index, 0)

  516.         transition1_bias = P.Gather()(self.transition1_biases, index, 0)

  517.         transition2_weight = P.Gather()(self.transition2_weights, index, 0)

  518.         transition2_bias = P.Gather()(self.transition2_biases, index, 0)

  519.         act, _, _ = self.input_layer_norm(act.astype(mstype.float32), input_layer_norm_gamma.astype(mstype.float32),

  520.                                           input_layer_norm_beta.astype(mstype.float32))

  521.         if self.slice_num:

  522.             act_ori_shape = P.Shape()(act)

  523.             slice_shape = (self.slice_num, -1) + act_ori_shape[1:]

  524.             act = P.Reshape()(act, slice_shape).astype(mstype.float16)

  525. 

  526.             slice_idx = 0

  527.             slice_idx_tensor = self.idx

  528.             act_tuple = ()

  529. 

  530.             act_slice = P.Gather()(act, slice_idx_tensor, 0)

  531.             act_shape = P.Shape()(act_slice)

  532.             if len(act_shape) != 2:

  533.                 act_slice = P.Reshape()(act_slice, (-1, act_shape[-1]))

  534.             act_slice = self.relu(

  535.                 P.BiasAdd()(self.matmul(act_slice.astype(mstype.float16), transition1_weight.astype(mstype.float16)),

  536.                             transition1_bias.astype(mstype.float16)).astype(mstype.float32))

  537.             act_slice = P.BiasAdd()(

  538.                 self.matmul(act_slice.astype(mstype.float16), transition2_weight.astype(mstype.float16)),

  539.                 transition2_bias.astype(mstype.float16))

  540.             act_slice = P.Reshape()(act_slice, act_shape)

  541.             act_slice = P.Reshape()(act_slice, ((1,) + P.Shape()(act_slice)))

  542.             act_tuple = act_tuple + (act_slice,)

  543.             slice_idx += 1

  544.             slice_idx_tensor += 1

  545. 

  546.             while slice_idx < self.slice_num:

  547.                 act_slice = P.Gather()(act, slice_idx_tensor, 0)

  548.                 act_slice = F.depend(act_slice, act_tuple[-1])

  549.                 act_shape = P.Shape()(act_slice)

  550.                 if len(act_shape) != 2:

  551.                     act_slice = P.Reshape()(act_slice, (-1, act_shape[-1]))

  552.                 act_slice = self.relu(P.BiasAdd()(

  553.                     self.matmul(act_slice.astype(mstype.float16), transition1_weight.astype(mstype.float16)),

  554.                     transition1_bias.astype(mstype.float16)).astype(mstype.float32))

  555.                 act_slice = P.BiasAdd()(

  556.                     self.matmul(act_slice.astype(mstype.float16), transition2_weight.astype(mstype.float16)),

  557.                     transition2_bias.astype(mstype.float16))

  558.                 act_slice = P.Reshape()(act_slice, act_shape)

  559.                 act_slice = P.Reshape()(act_slice, ((1,) + P.Shape()(act_slice)))

  560.                 act_tuple = act_tuple + (act_slice,)

  561.                 slice_idx += 1

  562.                 slice_idx_tensor += 1

  563. 

  564.             act = P.Concat()(act_tuple)

  565.             act = P.Reshape()(act, act_ori_shape)

  566.             return act

  567. 

  568.         act_shape = P.Shape()(act)

  569.         if len(act_shape) != 2:

  570.             act = P.Reshape()(act, (-1, act_shape[-1]))

  571.         act = self.relu(

  572.             P.BiasAdd()(self.matmul(act.astype(mstype.float16), transition1_weight.astype(mstype.float16)),

  573.                         transition1_bias.astype(mstype.float16)).astype(mstype.float32))

  574.         act = P.BiasAdd()(self.matmul(act.astype(mstype.float16), transition2_weight.astype(mstype.float16)),

  575.                           transition2_bias.astype(mstype.float16))

  576.         act = P.Reshape()(act, act_shape)

  577.         return act

  578. 

  579. 

  580. class OuterProductMean(nn.Cell):

  581.     '''outerproduct mean'''

  582.     def __init__(self, config, act_dim, num_output_channel, batch_size=None, slice_num=0):

  583.         super(OuterProductMean, self).__init__()

  584.         self.num_output_channel = num_output_channel

  585.         self.config = config

  586.         self.layer_norm_input = P.LayerNorm(begin_norm_axis=-1, begin_params_axis=-1, epsilon=1e-5)

  587.         self.matmul_trans_b = P.MatMul(transpose_b=True)

  588.         self.matmul = P.MatMul()

  589.         self.batch_matmul_trans_b = P.BatchMatMul(transpose_b=True)

  590.         self.act_dim = act_dim

  591.         self.batch_size = batch_size

  592.         self.slice_num = slice_num

  593.         self.idx = Tensor(0, mstype.int32)

  594.         self._init_parameter()

  595. 

  596.     def _init_parameter(self):

  597.         '''init parameter'''

  598.         self.layer_norm_input_gammas = Parameter(Tensor(np.zeros((self.batch_size, self.act_dim)), mstype.float32))

  599.         self.layer_norm_input_betas = Parameter(Tensor(np.zeros((self.batch_size, self.act_dim)), mstype.float32))

  600.         self.left_projection_weights = Parameter(

  601.             Tensor(np.zeros((self.batch_size, self.config.num_outer_channel, self.act_dim)), mstype.float32))

  602.         self.left_projection_biases = Parameter(

  603.             Tensor(np.zeros((self.batch_size, self.config.num_outer_channel)), mstype.float32))

  604.         self.right_projection_weights = Parameter(

  605.             Tensor(np.zeros((self.batch_size, self.config.num_outer_channel, self.act_dim)), mstype.float32))

  606.         self.right_projection_biases = Parameter(

  607.             Tensor(np.zeros((self.batch_size, self.config.num_outer_channel)), mstype.float32))

  608.         self.linear_output_weights = Parameter(Tensor(np.zeros(

  609.             (self.batch_size, self.num_output_channel, self.config.num_outer_channel *

  610.              self.config.num_outer_channel)), mstype.float32))

  611.         self.o_biases = Parameter(Tensor(np.zeros((self.batch_size, self.num_output_channel)), mstype.float32))

  612. 

  613.     def construct(self, act, mask, index):

  614.         '''construct'''

  615.         layer_norm_input_gamma = P.Gather()(self.layer_norm_input_gammas, index, 0)

  616.         layer_norm_input_beta = P.Gather()(self.layer_norm_input_betas, index, 0)

  617.         left_projection_weight = P.Gather()(self.left_projection_weights, index, 0)

  618.         left_projection_bias = P.Gather()(self.left_projection_biases, index, 0)

  619.         right_projection_weight = P.Gather()(self.right_projection_weights, index, 0)

  620.         right_projection_bias = P.Gather()(self.right_projection_biases, index, 0)

  621.         linear_output_weight = P.Gather()(self.linear_output_weights, index, 0)

  622.         linear_output_bias = P.Gather()(self.o_biases, index, 0)

  623.         mask = mask[..., None]

  624.         act, _, _ = self.layer_norm_input(act.astype(mstype.float32),

  625.                                           layer_norm_input_gamma.astype(mstype.float32),

  626.                                           layer_norm_input_beta.astype(mstype.float32))

  627.         act_shape = P.Shape()(act)

  628.         if len(act_shape) != 2:

  629.             act = P.Reshape()(act, (-1, act_shape[-1]))

  630.         out_shape = act_shape[:-1] + (-1,)

  631.         left_act = mask * P.Reshape()(P.BiasAdd()(self.matmul_trans_b(act.astype(mstype.float16),

  632.                                                                       left_projection_weight.astype(mstype.float16)),

  633.                                                   left_projection_bias.astype(mstype.float16)), out_shape)

  634.         right_act = mask * P.Reshape()(P.BiasAdd()(self.matmul_trans_b(act.astype(mstype.float16),

  635.                                                                        right_projection_weight.astype(mstype.float16)),

  636.                                                    right_projection_bias.astype(mstype.float16)), out_shape)

  637.         _, d, e = right_act.shape

  638.         if self.slice_num:

  639.             left_act_shape = P.Shape()(left_act)

  640.             slice_shape = (left_act_shape[0],) + (self.slice_num, -1) + (left_act_shape[-1],)

  641.             left_act = P.Reshape()(left_act, slice_shape)

  642.             slice_idx = 0

  643.             slice_idx_tensor = self.idx

  644.             act_tuple = ()

  645.             left_act_slice = P.Gather()(left_act, slice_idx_tensor, 1)

  646.             a, b, c = left_act_slice.shape

  647.             left_act_slice = P.Reshape()(mnp.transpose(left_act_slice.astype(mstype.float16), [2, 1, 0]), (-1, a))

  648.             right_act = P.Reshape()(right_act, (a, -1))

  649.             act_slice = P.Reshape()(P.Transpose()(P.Reshape()(self.matmul(left_act_slice.astype(mstype.float16),

  650.                                                                           right_act.astype(mstype.float16)),

  651.                                                               (c, b, d, e)), (2, 1, 0, 3)), (d, b, c * e))

  652.             act_slice_shape = P.Shape()(act_slice)

  653.             if len(act_shape) != 2:

  654.                 act_slice = P.Reshape()(act_slice, (-1, act_slice_shape[-1]))

  655.             act_slice = P.Reshape()(P.BiasAdd()(self.matmul_trans_b(act_slice.astype(mstype.float16),

  656.                                                                     linear_output_weight.astype(mstype.float16)),

  657.                                                 linear_output_bias.astype(mstype.float16)), (d, b, -1))

  658.             act_slice = mnp.transpose(act_slice.astype(mstype.float16), [1, 0, 2])

  659.             act_slice = P.Reshape()(act_slice, ((1,) + P.Shape()(act_slice)))

  660.             act_tuple = act_tuple + (act_slice,)

  661.             slice_idx += 1

  662.             slice_idx_tensor += 1

  663.             while slice_idx < self.slice_num:

  664.                 left_act_slice = P.Gather()(left_act, slice_idx_tensor, 1)

  665.                 left_act_slice = F.depend(left_act_slice, act_tuple[-1])

  666.                 a, b, c = left_act_slice.shape

  667.                 left_act_slice = P.Reshape()(mnp.transpose(left_act_slice.astype(mstype.float16), [2, 1, 0]), (-1, a))

  668.                 right_act = P.Reshape()(right_act, (a, -1))

  669.                 act_slice = P.Reshape()(P.Transpose()(P.Reshape()(self.matmul(left_act_slice.astype(mstype.float16),

  670.                                                                               right_act.astype(mstype.float16)),

  671.                                                                   (c, b, d, e)), (2, 1, 0, 3)), (d, b, c * e))

  672.                 act_slice_shape = P.Shape()(act_slice)

  673.                 if len(act_shape) != 2:

  674.                     act_slice = P.Reshape()(act_slice, (-1, act_slice_shape[-1]))

  675.                 act_slice = P.Reshape()(P.BiasAdd()(self.matmul_trans_b(act_slice.astype(mstype.float16),

  676.                                                                         linear_output_weight.astype(mstype.float16)),

  677.                                                     linear_output_bias.astype(mstype.float16)), (d, b, -1))

  678.                 act_slice = mnp.transpose(act_slice.astype(mstype.float16), [1, 0, 2])

  679.                 act_slice = P.Reshape()(act_slice, ((1,) + P.Shape()(act_slice)))

  680.                 act_tuple = act_tuple + (act_slice,)

  681.                 slice_idx += 1

  682.                 slice_idx_tensor += 1

  683.             act = P.Concat()(act_tuple)

  684.             act_shape = P.Shape()(act)

  685.             act = P.Reshape()(act, (-1, act_shape[-2], act_shape[-1]))

  686.             epsilon = 1e-3

  687.             tmp_mask = P.Transpose()(mask.astype(mstype.float16), (2, 1, 0))

  688.             norm = P.Transpose()(self.batch_matmul_trans_b(tmp_mask.astype(mstype.float16),

  689.                                                            tmp_mask.astype(mstype.float16)),

  690.                                  (1, 2, 0)).astype(mstype.float32)

  691.             act /= epsilon + norm

  692.             return act

  693. 

  694.         a, b, c = left_act.shape

  695.         left_act = P.Reshape()(mnp.transpose(left_act.astype(mstype.float16), [2, 1, 0]), (-1, a))

  696.         right_act = P.Reshape()(right_act, (a, -1))

  697.         act = P.Reshape()(P.Transpose()(P.Reshape()(self.matmul(left_act.astype(mstype.float16),

  698.                                                                 right_act.astype(mstype.float16)),

  699.                                                     (c, b, d, e)), (2, 1, 0, 3)), (d, b, c * e))

  700.         act_shape = P.Shape()(act)

  701.         if len(act_shape) != 2:

  702.             act = P.Reshape()(act, (-1, act_shape[-1]))

  703.         act = P.Reshape()(P.BiasAdd()(self.matmul_trans_b(act.astype(mstype.float16),

  704.                                                           linear_output_weight.astype(mstype.float16)),

  705.                                       linear_output_bias.astype(mstype.float16)), (d, b, -1))

  706.         act = mnp.transpose(act.astype(mstype.float16), [1, 0, 2])

  707.         epsilon = 1e-3

  708.         tmp_mask = P.Transpose()(mask.astype(mstype.float16), (2, 1, 0))

  709.         norm = P.Transpose()(self.batch_matmul_trans_b(tmp_mask.astype(mstype.float16),

  710.                                                        tmp_mask.astype(mstype.float16)),

  711.                              (1, 2, 0)).astype(mstype.float32)

  712.         act /= epsilon + norm

  713.         return act

  714. 

  715. class TriangleMultiplication(nn.Cell):

  716.     '''triangle multiplication'''

  717.     def __init__(self, config, layer_norm_dim, batch_size):

  718.         super(TriangleMultiplication, self).__init__()

  719.         self.config = config

  720.         self.layer_norm = P.LayerNorm(begin_norm_axis=-1, begin_params_axis=-1, epsilon=1e-5)

  721.         self.matmul = P.MatMul(transpose_b=True)

  722.         self.sigmoid = nn.Sigmoid()

  723.         self.batch_matmul_trans_b = P.BatchMatMul(transpose_b=True)

  724.         equation = ["ikc,jkc->ijc", "kjc,kic->ijc"]

  725.         if self.config.equation not in equation:

  726.             print("TriangleMultiplication Not Suppl")

  727.         if self.config.equation == "ikc,jkc->ijc":

  728.             self.equation = True

  729.         elif self.config.equation == "kjc,kic->ijc":

  730.             self.equation = False

  731.         else:

  732.             self.equation = None

  733.         self.batch_size = batch_size

  734.         self.layer_norm_dim = layer_norm_dim

  735.         self._init_parameter()

  736. 

  737.     def _init_parameter(self):

  738.         '''init parameter'''

  739.         self.layer_norm_input_gammas = Parameter(

  740.             Tensor(np.zeros((self.batch_size, self.layer_norm_dim)), mstype.float32))

  741.         self.layer_norm_input_betas = Parameter(

  742.             Tensor(np.zeros((self.batch_size, self.layer_norm_dim)), mstype.float32))

  743.         self.left_projection_weights = Parameter(

  744.             Tensor(np.zeros((self.batch_size, self.config.num_intermediate_channel, self.layer_norm_dim)),

  745.                    mstype.float32))

  746.         self.left_projection_biases = Parameter(

  747.             Tensor(np.zeros((self.batch_size, self.config.num_intermediate_channel)), mstype.float32))

  748.         self.right_projection_weights = Parameter(

  749.             Tensor(np.zeros((self.batch_size, self.config.num_intermediate_channel, self.layer_norm_dim)),

  750.                    mstype.float32))

  751.         self.right_projection_biases = Parameter(

  752.             Tensor(np.zeros((self.batch_size, self.config.num_intermediate_channel)), mstype.float32))

  753.         self.left_gate_weights = Parameter(

  754.             Tensor(np.zeros((self.batch_size, self.config.num_intermediate_channel, self.layer_norm_dim)),

  755.                    mstype.float32))

  756.         self.left_gate_biases = Parameter(

  757.             Tensor(np.zeros((self.batch_size, self.config.num_intermediate_channel)), mstype.float32))

  758.         self.right_gate_weights = Parameter(

  759.             Tensor(np.zeros((self.batch_size, self.config.num_intermediate_channel, self.layer_norm_dim)),

  760.                    mstype.float32))

  761.         self.right_gate_biases = Parameter(

  762.             Tensor(np.zeros((self.batch_size, self.config.num_intermediate_channel)), mstype.float32))

  763.         self.center_layer_norm_gammas = Parameter(

  764.             Tensor(np.zeros((self.batch_size, self.layer_norm_dim)), mstype.float32))

  765.         self.center_layer_norm_betas = Parameter(

  766.             Tensor(np.zeros((self.batch_size, self.layer_norm_dim)), mstype.float32))

  767.         self.output_projection_weights = Parameter(

  768.             Tensor(np.zeros((self.batch_size, self.layer_norm_dim, self.layer_norm_dim)), mstype.float32))

  769.         self.output_projection_biases = Parameter(

  770.             Tensor(np.zeros((self.batch_size, self.layer_norm_dim)), mstype.float32))

  771.         self.gating_linear_weights = Parameter(

  772.             Tensor(np.zeros((self.batch_size, self.layer_norm_dim, self.layer_norm_dim)), mstype.float32))

  773.         self.gating_linear_biases = Parameter(Tensor(np.zeros((self.batch_size, self.layer_norm_dim)), mstype.float32))

  774. 

  775.     def construct(self, act, mask, index):

  776.         '''construct'''

  777.         layer_norm_input_gamma = P.Gather()(self.layer_norm_input_gammas, index, 0)

  778.         layer_norm_input_beta = P.Gather()(self.layer_norm_input_betas, index, 0)

  779.         left_projection_weight = P.Gather()(self.left_projection_weights, index, 0)

  780.         left_projection_bias = P.Gather()(self.left_projection_biases, index, 0)

  781.         right_projection_weight = P.Gather()(self.right_projection_weights, index, 0)

  782.         right_projection_bias = P.Gather()(self.right_projection_biases, index, 0)

  783.         left_gate_weight = P.Gather()(self.left_gate_weights, index, 0)

  784.         left_gate_bias = P.Gather()(self.left_gate_biases, index, 0)

  785.         right_gate_weight = P.Gather()(self.right_gate_weights, index, 0)

  786.         right_gate_bias = P.Gather()(self.right_gate_biases, index, 0)

  787.         center_layer_norm_gamma = P.Gather()(self.center_layer_norm_gammas, index, 0)

  788.         center_layer_norm_beta = P.Gather()(self.center_layer_norm_betas, index, 0)

  789.         output_projection_weight = P.Gather()(self.output_projection_weights, index, 0)

  790.         output_projection_bias = P.Gather()(self.output_projection_biases, index, 0)

  791.         gating_linear_weight = P.Gather()(self.gating_linear_weights, index, 0)

  792.         gating_linear_bias = P.Gather()(self.gating_linear_biases, index, 0)

  793. 

  794.         mask = mask[..., None]

  795.         act, _, _ = self.layer_norm(act.astype(mstype.float32),

  796.                                     layer_norm_input_gamma.astype(mstype.float32),

  797.                                     layer_norm_input_beta.astype(mstype.float32))

  798.         input_act = act

  799.         act_shape = P.Shape()(act)

  800.         if len(act_shape) != 2:

  801.             act = P.Reshape()(act, (-1, act_shape[-1]))

  802.         out_shape = act_shape[:-1] + (-1,)

  803.         left_projection = mask * P.Reshape()(

  804.             P.BiasAdd()(self.matmul(act.astype(mstype.float16), left_projection_weight.astype(mstype.float16)),

  805.                         left_projection_bias.astype(mstype.float16)), out_shape)

  806.         left_projection = left_projection.astype(mstype.float16)

  807.         act = F.depend(act, left_projection)

  808. 

  809.         left_gate_values = self.sigmoid(P.Reshape()(

  810.             P.BiasAdd()(self.matmul(act.astype(mstype.float16), left_gate_weight.astype(mstype.float16)),

  811.                         left_gate_bias.astype(mstype.float16)), out_shape).astype(mstype.float32))

  812.         left_proj_act = left_projection * left_gate_values

  813.         act = F.depend(act, left_proj_act)

  814. 

  815.         right_projection = mask * P.Reshape()(

  816.             P.BiasAdd()(self.matmul(act.astype(mstype.float16), right_projection_weight.astype(mstype.float16)),

  817.                         right_projection_bias.astype(mstype.float16)), out_shape)

  818.         right_projection = right_projection.astype(mstype.float16)

  819.         act = F.depend(act, right_projection)

  820. 

  821.         right_gate_values = self.sigmoid(P.Reshape()(

  822.             P.BiasAdd()(self.matmul(act.astype(mstype.float16), right_gate_weight.astype(mstype.float16)),

  823.                         right_gate_bias.astype(mstype.float16)), out_shape).astype(mstype.float32))

  824.         right_proj_act = right_projection * right_gate_values

  825.         left_proj_act = F.depend(left_proj_act, right_proj_act)

  826. 

  827.         if self.equation is not None:

  828.             if self.equation:

  829.                 left_proj_act_tmp = P.Transpose()(left_proj_act.astype(mstype.float16), (2, 0, 1))

  830.                 right_proj_act_tmp = P.Transpose()(right_proj_act.astype(mstype.float16), (2, 0, 1))

  831.                 act = self.batch_matmul_trans_b(left_proj_act_tmp, right_proj_act_tmp)

  832.                 act = P.Transpose()(act, (1, 2, 0)).astype(mstype.float32)

  833.             else:

  834.                 left_proj_act_tmp = P.Transpose()(left_proj_act.astype(mstype.float16), (2, 1, 0))

  835.                 right_proj_act_tmp = P.Transpose()(right_proj_act.astype(mstype.float16), (2, 1, 0))

  836.                 act = self.batch_matmul_trans_b(left_proj_act_tmp, right_proj_act_tmp)

  837.                 act = P.Transpose()(act, (2, 1, 0)).astype(mstype.float32)

  838.         act, _, _ = self.layer_norm(act.astype(mstype.float32),

  839.                                     center_layer_norm_gamma.astype(mstype.float32),

  840.                                     center_layer_norm_beta.astype(mstype.float32))

  841.         act_shape = P.Shape()(act)

  842.         if len(act_shape) != 2:

  843.             act = P.Reshape()(act, (-1, act_shape[-1]))

  844.         out_shape = act_shape[:-1] + (-1,)

  845.         act = P.Reshape()(

  846.             P.BiasAdd()(self.matmul(act.astype(mstype.float16), output_projection_weight.astype(mstype.float16)),

  847.                         output_projection_bias.astype(mstype.float16)), out_shape)

  848.         input_act_shape = P.Shape()(input_act)

  849.         if len(input_act_shape) != 2:

  850.             input_act = P.Reshape()(input_act, (-1, input_act_shape[-1]))

  851.         out_shape = input_act_shape[:-1] + (-1,)

  852.         gate_values = self.sigmoid(P.Reshape()(

  853.             P.BiasAdd()(self.matmul(input_act.astype(mstype.float16), gating_linear_weight.astype(mstype.float16)),

  854.                         gating_linear_bias.astype(mstype.float16)), out_shape).astype(mstype.float32))

  855.         act = act * gate_values

  856.         return act

  857. 

  858. 

  859. class TriangleAttention(nn.Cell):

  860.     '''triangle attention'''

  861.     def __init__(self, config, layer_norm_dim, batch_size=None, slice_num=0):

  862.         super(TriangleAttention, self).__init__()

  863.         self.config = config

  864.         self.orientation_is_per_column = (self.config.orientation == 'per_column')

  865.         self.init_factor = Tensor(1. / np.sqrt(layer_norm_dim), mstype.float32)

  866.         self.query_norm = P.LayerNorm(begin_norm_axis=-1, begin_params_axis=-1, epsilon=1e-5)

  867.         self.matmul = P.MatMul(transpose_b=True)

  868.         self.attn_mod = Attention(self.config, layer_norm_dim, layer_norm_dim, layer_norm_dim, batch_size)

  869.         self.batch_size = batch_size

  870.         self.slice_num = slice_num

  871.         self.layer_norm_dim = layer_norm_dim

  872.         self.idx = Tensor(0, mstype.int32)

  873.         self._init_parameter()

  874. 

  875.     def _init_parameter(self):

  876.         '''init parameter'''

  877.         self.query_norm_gammas = Parameter(Tensor(np.zeros((self.batch_size, self.layer_norm_dim)), mstype.float32))

  878.         self.query_norm_betas = Parameter(Tensor(np.zeros((self.batch_size, self.layer_norm_dim)), mstype.float32))

  879.         self.feat_2d_weights = Parameter(

  880.             Tensor(np.zeros((self.batch_size, self.config.num_head, self.layer_norm_dim)), mstype.float32))

  881. 

  882.     def construct(self, pair_act, pair_mask, index):

  883.         '''construct'''

  884.         query_norm_gamma = P.Gather()(self.query_norm_gammas, index, 0)

  885.         query_norm_beta = P.Gather()(self.query_norm_betas, index, 0)

  886.         feat_2d_weight = P.Gather()(self.feat_2d_weights, index, 0)

  887.         if self.orientation_is_per_column:

  888.             pair_act = mnp.swapaxes(pair_act, -2, -3)

  889.             pair_mask = mnp.swapaxes(pair_mask, -1, -2)

  890.         bias = (1e9 * (pair_mask - 1.))[:, None, None, :]

  891.         pair_act, _, _ = self.query_norm(pair_act.astype(mstype.float32),

  892.                                          query_norm_gamma.astype(mstype.float32),

  893.                                          query_norm_beta.astype(mstype.float32))

  894.         q, k, _ = pair_act.shape

  895.         nonbatched_bias = self.matmul(P.Reshape()(pair_act.astype(mstype.float16), (-1, pair_act.shape[-1])),

  896.                                       feat_2d_weight.astype(mstype.float16))

  897.         nonbatched_bias = P.Transpose()(P.Reshape()(nonbatched_bias, (q, k, -1)), (2, 0, 1))

  898.         if self.slice_num:

  899.             pair_act_ori_shape = P.Shape()(pair_act)

  900.             slice_shape = (self.slice_num, -1) + pair_act_ori_shape[1:]

  901.             pair_act = P.Reshape()(pair_act, slice_shape).astype(mstype.float16)

  902.             bias_shape = P.Shape()(bias)

  903.             bias = P.Reshape()(bias, slice_shape[:2] + bias_shape[1:])

  904. 

  905.             slice_idx = 0

  906.             slice_idx_tensor = self.idx

  907.             pair_act_tuple = ()

  908. 

  909.             pair_act_slice = P.Gather()(pair_act, slice_idx_tensor, 0)

  910.             bias_slice = P.Gather()(bias, slice_idx_tensor, 0)

  911.             pair_act_slice = self.attn_mod(pair_act_slice, pair_act_slice, bias_slice, index, nonbatched_bias)

  912.             pair_act_slice = P.Reshape()(pair_act_slice, ((1,) + P.Shape()(pair_act_slice)))

  913.             pair_act_tuple = pair_act_tuple + (pair_act_slice,)

  914.             slice_idx += 1

  915.             slice_idx_tensor += 1

  916. 

  917.             while slice_idx < self.slice_num:

  918.                 pair_act_slice = P.Gather()(pair_act, slice_idx_tensor, 0)

  919.                 pair_act_slice = F.depend(pair_act_slice, pair_act_tuple[-1])

  920.                 bias_slice = P.Gather()(bias, slice_idx_tensor, 0)

  921.                 pair_act_slice = self.attn_mod(pair_act_slice, pair_act_slice, bias_slice, index, nonbatched_bias)

  922.                 pair_act_slice = P.Reshape()(pair_act_slice, ((1,) + P.Shape()(pair_act_slice)))

  923.                 pair_act_tuple = pair_act_tuple + (pair_act_slice,)

  924.                 slice_idx += 1

  925.                 slice_idx_tensor += 1

  926.             pair_act = P.Concat()(pair_act_tuple)

  927.             pair_act = P.Reshape()(pair_act, pair_act_ori_shape)

  928. 

  929.             if self.orientation_is_per_column:

  930.                 pair_act = mnp.swapaxes(pair_act, -2, -3)

  931.             return pair_act

  932. 

  933.         pair_act = self.attn_mod(pair_act, pair_act, bias, index, nonbatched_bias)

  934.         if self.orientation_is_per_column:

  935.             pair_act = mnp.swapaxes(pair_act, -2, -3)

  936.         return pair_act