torchcell.models.deep_set.DeepSet

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

Bases: Module

Permutation-invariant network applying per-node MLPs then set aggregation.

node_layers_forward(x: Tensor) → Tensor

Process node features through node layers.

set_layers_forward(x_aggregated: Tensor) → Tensor

Process aggregated features through set layers.

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

Apply node layers, aggregate per set, then apply set layers.

Parameters:

• x – Node feature tensor of shape (num_nodes, in_channels).

• batch – Set-assignment index for each node.

Returns:

A tuple of per-node features and per-set output features.

Raises:

ValueError – self.aggregation is not sum or mean.