emmi.modules.attention.irregular_nat_attention¶

Classes¶

IrregularNatAttention

Implementation the attention implemenation of Neighbourhood Attention Transformer (NAT) for irregular grinds .

Module Contents¶

class emmi.modules.attention.irregular_nat_attention.IrregularNatAttention(config)¶

Bases: torch.nn.Module

Implementation the attention implemenation of Neighbourhood Attention Transformer (NAT) for irregular grinds . - Source: https://arxiv.org/abs/2204.07143

Initialize the IrregularNatAttention module.

Parameters:: config (emmi.schemas.modules.attention.AttentionConfig) – configuration of the attention module.

radius = None¶

max_degree = None¶

num_heads = None¶

head_dim¶

relative_position_bias_mlp¶

qkv¶

proj¶

reset_parameters()¶

Restet the parameters of the PerceiverAttention module with a specific initialization method.

Raises:: NotImplementedError – when a specific initialization method is not implemented. Either use “torch”, or “truncnormal002”.
Return type:: None

forward(x, pos)¶

Forward function of the IrregularNatAttention module.

Parameters:

x (torch.Tensor) – input tensor, shape (batch size, number of points/tokens, hidden_dim).
pos (torch.Tensor) – coordinates points for radius nearest neighbors.

Raises:

NotImplementedError – Only supports 3D input.

Returns:

Returns the output of the attention module.

Return type:

torch.Tensor