# Auto-extracted class source (static)

class GNNDiracRRF(nn.Module):
    def __init__(self, in_dim: int, hidden_dim: int, out_dim: int, num_layers: int, z_dim: int,
                 alpha_attn: float = 1.0, dropout: float = 0.1):
        super().__init__()
        self.z_dim = z_dim
        self.layers = nn.ModuleList()
        self.layers.append(DiracGraphConv(in_dim, hidden_dim, alpha=alpha_attn))
        for _ in range(num_layers - 2):
            self.layers.append(DiracGraphConv(hidden_dim, hidden_dim, alpha=alpha_attn))
        self.layers.append(DiracGraphConv(hidden_dim, out_dim, alpha=alpha_attn))
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
        h = x
        for i, layer in enumerate(self.layers):
            h = layer(h, edge_index, z)
            if i < len(self.layers) - 1:
                h = F.gelu(h)
                h = self.dropout(h)
        return h