code: def multi_head_attention(
      queries,
      keys,
      values,
      attn_bias,
      d_key,
      d_value,
      d_model,
      input_mask,
      n_head=1,
      dropout_rate=0.0,
      cache=None,
      param_initializer=None,
      name="multi_head_att",
  ):
      """
      Multi-Head Attention. Note that attn_bias is added to the logit before
      computing softmax activiation to mask certain selected positions so that
      they will not considered in attention weights.
      """
      keys = queries if keys is None else keys
      values = keys if values is None else values
  
      def __compute_qkv(queries, keys, values, n_head, d_key, d_value):
          """
          Add linear projection to queries, keys, and values.
          """
          q = layers.fc(
              input=queries,
              size=d_key * n_head,
              num_flatten_dims=len(queries.shape) - 1,
              param_attr=fluid.ParamAttr(
                  name=name + "_query_fc.w_0",
                  initializer=param_initializer,
              ),
              bias_attr=name + "_query_fc.b_0",
          )
          k = layers.fc(
              input=keys,
              size=d_key * n_head,
              num_flatten_dims=len(keys.shape) - 1,
              param_attr=fluid.ParamAttr(
                  name=name + "_key_fc.w_0",
                  initializer=param_initializer,
              ),
              bias_attr=name + "_key_fc.b_0",
          )
          v = layers.fc(
              input=values,
              size=d_value * n_head,
              num_flatten_dims=len(values.shape) - 1,
              param_attr=fluid.ParamAttr(
                  name=name + "_value_fc.w_0",
                  initializer=param_initializer,
              ),
              bias_attr=name + "_value_fc.b_0",
          )
          return q, k, v
  
      def __split_heads(x, n_head):
          """
          Reshape the last dimension of inpunt tensor x so that it becomes two
          dimensions and then transpose. Specifically, input a tensor with shape
          [bs, max_sequence_length, n_head * hidden_dim] then output a tensor
          with shape [bs, n_head, max_sequence_length, hidden_dim].
          """
          hidden_size = x.shape[-1]
          # The value 0 in shape attr means copying the corresponding dimension
          # size of the input as the output dimension size.
          reshaped = layers.reshape(
              x=x,
              shape=[0, 0, n_head, hidden_size // n_head],
              inplace=True,
          )
  
          # permuate the dimensions into:
          # [batch_size, n_head, max_sequence_len, hidden_size_per_head]
          return layers.transpose(x=reshaped, perm=[0, 2, 1, 3])
  
      def __combine_heads(x):
          """
          Transpose and then reshape the last two dimensions of inpunt tensor x
          so that it becomes one dimension, which is reverse to __split_heads.
          """
          if len(x.shape) == 3:
              return x
          if len(x.shape) != 4:
              raise ValueError("Input(x) should be a 4-D Tensor.")
          trans_x = layers.transpose(x, perm=[0, 2, 1, 3])
          # The value 0 in shape attr means copying the corresponding dimension
          # size of the input as the output dimension size.
          # trans_x.desc.set_shape((-1, 1, n_head, d_value))
          return layers.reshape(
              x=trans_x, shape=[0, 0, d_model], inplace=True
          )
  
      q, k, v = __compute_qkv(
          queries, keys, values, n_head, d_key, d_value
      )
      q = to_3d(q)
      k = to_3d(k)
      v = to_3d(v)
  
      if cache is not None:  # use cache and concat time steps
          # Since the inplace reshape in __split_heads changes the shape of k and
          # v, which is the cache input for next time step, reshape the cache
          # input from the previous time step first.
          k = cache["k"] = layers.concat(
              [layers.reshape(cache["k"], shape=[0, 0, d_model]), k],
              axis=1,
          )
          v = cache["v"] = layers.concat(
              [layers.reshape(cache["v"], shape=[0, 0, d_model]), v],
              axis=1,
          )
  
      out, _ = sparse_scaled_dot_product_attention(
          q, k, v, input_mask, dropout_rate, n_head, d_key, d_value
      )
  
      out = to_2d(out)
  
      # Project back to the model size.
      proj_out = layers.fc(
          input=out,
          size=d_model,
          num_flatten_dims=len(out.shape) - 1,
          param_attr=fluid.ParamAttr(
              name=name + "_output_fc.w_0",
              initializer=param_initializer,
          ),
          bias_attr=name + "_output_fc.b_0",
      )
      return proj_out, _
  
file_path: PGL/legacy/examples/xformer/transformer_encoder_sparse.py
imports: from functools import partial
  import paddle.fluid
  import paddle.fluid.layers
  import paddle.fluid.layers
  from .sparse_scaled_dot_product_attention import sparse_scaled_dot_product_attention
label_desc: Use image node_id 1 for calling a global function with example usage: multi_head_attention(queries, keys, values, attn_bias, d_key, d_value, d_model, input_mask, n_head, dropout_rate, cache, param_initializer, name) and returns: proj_out,  _, q,  k,  v, layers, layers, x