grouped_vector_attention

GroupedVectorAttention(
    channels: int,
    num_groups: int,
    attn_drop: float = 0.0,
    qkv_bias: bool = True,
    pe_multiplier: bool = False,
    pe_bias: bool = True,
    norm: Union[str, Callable, None] = "batch_norm",
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_kwargs: Optional[Dict[str, Any]] = None,
    **kwargs: Unpack[MessagePassingParams],
)

Bases: MessagePassing

Vector attention over a neighborhood graph, with one weight vector shared by each group of channels.

The relation between a query and its neighbor keys is optionally scaled and shifted by a learned encoding of their relative position, then mapped to num_groups weights and softmax-normalized over each destination's neighbors.

edge_index lists [source, target] pairs: keys and values come from the source points and queries from the target points, so a bipartite graph is given as (x_source, x_target) and (pos_source, pos_target) pairs.

Parameters:

• channels (int) –

  Feature width of the queries, keys and values.

• num_groups (int) –

  Number of channel groups sharing one attention weight.

• attn_drop (float, default: 0.0 ) –

  Dropout probability on the attention weights.

• qkv_bias (bool, default: True ) –

  Whether the query, key and value projections carry a bias.

• pe_multiplier (bool, default: False ) –

  Whether to scale the query-key relation by an encoding of the relative position.

• pe_bias (bool, default: True ) –

  Whether to shift the query-key relation and the values by an encoding of the relative position.

• norm (Union[str, Callable, None], default: 'batch_norm' ) –

  Normalization type or callable.

• act (Union[str, Callable, None], default: 'relu' ) –

  Activation type or callable.

• act_kwargs (Optional[Dict[str, Any]], default: None ) –

  Extra activation arguments.

• norm_kwargs (Optional[Dict[str, Any]], default: None ) –

  Extra normalization arguments.

• **kwargs (Unpack[MessagePassingParams], default: {} ) –

  Extra arguments for MessagePassing.