lloca.backbone.transformer.Transformer

class lloca.backbone.transformer.Transformer(in_channels, attn_reps, out_channels, num_blocks, num_heads, checkpoint_blocks=False, attention_factor=1, mlp_factor=4, multi_query=False, dropout_prob=None, compile=False, compile_mode='default', compile_dynamic=True)

Bases: Module

Baseline LLoCa-Transformer.

Combines transformer blocks, each consisting of multi-head self-attention layers, an MLP, residual connections, and normalization layers.

Parameters:

• in_channels (int) – Number of input channels.

• attn_reps (str) – Representation of each attention head.

• out_channels (int) – Number of output channels.

• num_blocks (int) – Number of transformer blocks.

• num_heads (int) – Number of attention heads.

• checkpoint_blocks (bool) – Use gradient checkpointing for transformer blocks.

• attention_factor (int) – Factor by which the key, query, and value size is increased over the default value of hidden_channels / num_heads.

• mlp_factor (int) – Factor by which the activation size is increased over the default value of hidden_channels.

• multi_query (bool) – Use multi-query attention instead of multi-head attention.

• dropout_prob (float) – Dropout probability for output.

• compile (bool, optional) – Whether to compile the model with torch.compile, by default False.

• compile_mode (str) – torch.compile compilation mode, see torch docs for more information.

• compile_dynamic (bool) – Whether to use dynamic shapes with torch.compile, by default True.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(inputs, frames, **attn_kwargs)

Forward pass.

Parameters:

• inputs (Tensor) – Input data with shape (…, num_items, in_channels)

• frames (Frames) – Local frames used for invariant particle attention

• **attn_kwargs

Returns:

outputs – Outputs with shape (…, num_items, out_channels)

Return type:

Tensor