def multihead_fill(layer, qkv_same_embed_dim=True):
    """
    The default initialization of multi-head attention.

    Example:
        from paddleseg.cvlibs import param_init
        import paddle.nn as nn
        
        self_attn = nn.MultiHeadAttention(
            128, 8, dropout=False)
        param_init.multihead_fill(self_attn, True)
    """

    def _init_param_as_combined_linear_weight(p):
        bound = math.sqrt(6 / (3 * p.shape[0] + p.shape[1]))
        nn.initializer.Uniform(low=-bound, high=bound)(p)

    if qkv_same_embed_dim:
        _init_param_as_combined_linear_weight(layer.q_proj.weight)
        _init_param_as_combined_linear_weight(layer.k_proj.weight)
        _init_param_as_combined_linear_weight(layer.v_proj.weight)
        xavier_uniform(layer.out_proj.weight)
    else:
        for p in layer.parameters():
            if p.dim() > 1:
                xavier_uniform(p)

    The default initialization of multi-head attention.

    Example:
        from paddleseg.cvlibs import param_init
        import paddle.nn as nn
        
        self_attn = nn.MultiHeadAttention(
            128, 8, dropout=False)
        param_init.multihead_fill(self_attn, True)
    
paddleseg/cvlibs/param_init.py