torchcell.models.hetero_cell_bipartite_dango_gi_lazy.AttentionConvWrapper

class AttentionConvWrapper(conv: Module, target_dim: int, norm: str | None = None, activation: str | None = None, dropout: float = 0.1)

Bases: Module

Wrapper applying norm, activation, and dropout around a graph convolution.

forward(x: Tensor, edge_index: Tensor, edge_mask: Tensor | None = None, **kwargs: Any) → Tensor

Run the wrapped conv with optional edge masking, then norm/activation/dropout.

Parameters:

• x – Node features

• edge_index – Edge connectivity

• edge_mask – Optional edge mask for MaskedGINConv (True = keep edge)

• **kwargs – Additional arguments for the conv layer