torchcell.models.self_attention_deep_set.SelfAttentionDeepSet

class SelfAttentionDeepSet(in_channels: int, hidden_channels: int, out_channels: int, num_node_layers: int, num_set_layers: int, num_heads: int = 1, dropout_prob: float = 0.2, norm: str = 'batch', activation: str = 'relu', skip_node: bool = False, skip_set: bool = False)

Bases: Module

Deep Sets encoder with self-attention between node and set layers.

node_layers_forward(x: Tensor, batch: Tensor) → tuple[Tensor, list[list[Tensor]]]

Process node features through node layers.

set_layers_forward(x_summed: Tensor) → Tensor

Process aggregated features through set layers.

forward(x: Tensor, batch: Tensor) → tuple[Tensor, Tensor, list[list[Tensor]]]

Encode nodes, sum per graph, and return node, set, and attention.