kutacc_af2_triangle_multiplication_gate_and_out_linear

Intermediate step of triangle_multiplication, where gate is generated based on input_act through linear transformation and out is generated based on center_act through linear transformation.

Interface Definition

void kutacc_af2_triangle_multiplication_gate_and_out_linear(kutacc_tensor_h gate, kutacc_tensor_h out, kutacc_af2_tm_act_inputs_t *tm_acts_ptr, kutacc_tensor_h center_act,

kutacc_af2_tm_linear_weights_t *tm_weights_ptr, bool input_prepack);

Parameters

Table 1 Input parameter definition

Parameter	Type	Description	Input/Output
gate	kutacc_tensor_h	Gating tensor	Input
out	kutacc_tensor_h	Output tensor	Output
tm_acts_ptr	kutacc_af2_tm_act_inputs_t *	Pointer of the kutacc_af2_tm_act_inputs_t type. For details about the data structure, see Table 2.	Input
center_act	kutacc_tensor_h	Input obtained after permute and LayerNorm are performed on center_act in the equation step.	Input
tm_weights_ptr	kutacc_af2_tm_linear_weights_t *	Pointer of the kutacc_af2_tm_linear_weights_t type. For details about the data structure, see Table 3.	Input
input_prepack	bool	Where to perform prepack	Input

Table 2 kutacc_af2_tm_act_inputs_t structure definition

Parameter	Type	Description	Input/Output
n_res	int64_t	Number of residues	Input
n_res_gather	int64_t	Number of residues	Input
proj_act	kutacc_tensor_h	Left and right projections of act	Output
input_act	kutacc_tensor_h	Value generated after the input act is normalized by LayerNorm	Input
proj_act_gate	kutacc_tensor_h	Temporary gating tensor	Input

Table 3 kutacc_af2_tm_linear_weights_t structure definition

Parameter	Type	Description	Input/Output
c_o	int64_t	Output feature dimension	Input
c_i	int64_t	Input feature dimension	Input
gating_w	kutacc_tensor_h	Weight used for generating gate	Output
gating_b	kutacc_tensor_h	Bias used for generating gate	Input
output_proj_w	kutacc_tensor_h	Weight used for generating out	Input
output_proj_b	kutacc_tensor_h	Bias used for generating out	Input

Constraints on integer parameters of triangle_multiplication:

n_res, n_res_gather, c_o, c_i, c_z > 0

n_res * n_res_gather < INT64_MAX

c_i = c_o

In single-process mode, n_res must be equal to n_res_gather. In multi-process mode, this condition does not need to be met.

Parent topic: AlphaFold2