Nn¶

def forward(
    self, X: Float[torch.Tensor, "n_nodes dim"], A_hat: Float[torch.Tensor, "n_nodes n_nodes"] | None = None
) -> Float[torch.Tensor, "n_nodes dim"]:
    """Forward pass for the Kappa GCN layer.

    Args:
        X: Embedding matrix
        A_hat: Normalized adjacency matrix

    Returns:
        AXW: Transformed node features after message passing and nonlinear activation.
    """
    # 1. right-multiply X by W - mobius_matvec broadcasts correctly (verified)
    XW = self.manifold.manifold.mobius_matvec(m=self.W, x=X)

    # 2. left-multiply (X @ W) by A_hat - we need our own implementation for this
    if A_hat is None:
        AXW = XW
    elif isinstance(self.manifold, ProductManifold):
        XWs = self.manifold.factorize(XW)
        AXW = torch.hstack([self._left_multiply(A_hat, XW, M) for XW, M in zip(XWs, self.manifold.P, strict=False)])
    else:
        AXW = self._left_multiply(A_hat, XW, self.manifold)

    # 3. Apply nonlinearity - note that sigma is wrapped with our manifold.apply decorator
    AXW = self.sigma(AXW)

    return AXW

def forward(
    self, X: Float[torch.Tensor, "n_nodes dim"], A_hat: Float[torch.Tensor, "n_nodes n_nodes"] | None = None
) -> Float[torch.Tensor, "n_nodes out_dim"]:
    """Forward pass through all layers.

    Args:
        X: Input features
        A_hat: Adjacency matrix passed to each layer

    Returns:
        Output after passing through all layers
    """
    for layer in self.layers:
        X = layer(X, A_hat)
    return X

def append(self, layer: nn.Module) -> None:
    """Add a layer to the end of the sequence."""
    self.layers.append(layer)

def forward(
    self,
    X: Float[torch.Tensor, "n_nodes dim"],
    A_hat: Float[torch.Tensor, "n_nodes n_nodes"] | None = None,
    aggregate_logits: bool = False,
) -> Float[torch.Tensor, "n_nodes n_classes"]:
    """Forward pass through stereographic logits.

    Args:
        X: Input features
        A_hat: Optional adjacency matrix for logit aggregation
        aggregate_logits: Whether to aggregate logits using adjacency matrix

    Returns:
        Logits (or probabilities if apply_softmax=True)
    """
    # Compute logits based on manifold type
    if isinstance(self.manifold, ProductManifold):
        logits = self._get_logits_product_manifold(X, self.W, self.p_ks, self.manifold)
    else:
        logits = self._get_logits_single_manifold(X, self.W, self.p_ks, self.manifold, return_inner_products=False)

    # Optional aggregation via adjacency matrix
    if A_hat is not None and aggregate_logits:
        logits = A_hat @ logits

    # Optional softmax for classification
    if self.apply_softmax:
        logits = torch.softmax(logits, dim=-1)

    return logits

def forward(
    self,
    X: Float[torch.Tensor, "n_nodes dim"],
    A_hat: Float[torch.Tensor, "n_nodes n_nodes"] | None = None,
) -> Float[torch.Tensor, "n_nodes n_nodes"]:
    """Forward pass through Fermi-Dirac decoder.

    Args:
        X: Node embeddings
        A_hat: Ignored (for compatibility)
        return_pairwise: If True, return full pairwise matrix. If False, return flattened upper triangle.

    Returns:
        Edge probabilities (logits, apply sigmoid if needed)
    """
    # Compute pairwise distances
    pairwise_dist = self.manifold.pdist2(X)

    # Apply Fermi-Dirac transformation
    logits = -(pairwise_dist - self.bias) / self.temperature

    return logits

def forward(self, X: Float[torch.Tensor, "n_nodes dim"]) -> Float[torch.Tensor, "n_nodes dim"]:
    """Apply layer normalization on the stereographic manifold."""
    norm_X = self.stereographic_norm(X)
    output = geoopt.manifolds.stereographic.math.project(norm_X, self.curvatures)
    return output

def forward(
    self,
    Q: Float[torch.Tensor, "batch_size num_heads n_nodes head_dim"],
    K: Float[torch.Tensor, "batch_size num_heads n_nodes head_dim"],
    V: Float[torch.Tensor, "batch_size num_heads n_nodes head_dim"],
    mask: Float[torch.Tensor, "1 1 n_nodes n_nodes"],
) -> Float[torch.Tensor, "batch_size n_nodes dim"]:
    """Forward pass for the geometric linearized attention layer.

    Args:
        Q: Query tensor.
        K: Key tensor.
        V: Value tensor.
        mask: Mask tensor for attention.

    Returns:
        Output tensor after applying attention.
    """
    v1 = geoopt.manifolds.stereographic.math.parallel_transport0back(V, Q, k=self.curvatures)
    v2 = geoopt.manifolds.stereographic.math.parallel_transport0back(V, K, k=self.curvatures)

    gamma = geoopt.manifolds.stereographic.math.lambda_x(x=V, k=self.curvatures, keepdim=True, dim=-1)
    denominator = geoopt.utils.clamp_abs((gamma - 1), self._clamp_epsilon)

    x = ((gamma / denominator) * V) * mask[None, :, None]

    v1 = nn.functional.elu(v1) + 1
    v2 = (denominator * (nn.functional.elu(v2) + 1)) * mask[None, :, None]

    # Linearized approximation
    v2_cumsum = v2.sum(dim=-2)  # [B, H, D]
    D = torch.einsum("...nd,...d->...n", v1, v2_cumsum.type_as(v1))  # normalization terms
    D_inv = 1.0 / D.masked_fill_(D == 0, self._epsilon)
    context = torch.einsum("...nd,...ne->...de", v2, x)
    X = torch.einsum("...de,...nd,...n->...ne", context, v1, D_inv)

    X = geoopt.manifolds.stereographic.math.project(X, k=self.curvatures)
    X = geoopt.manifolds.stereographic.math.mobius_scalar_mul(
        torch.tensor(0.5, dtype=X.dtype, device=X.device), X, k=self.curvatures, dim=-1
    )
    X = geoopt.manifolds.stereographic.math.project(X, k=self.curvatures)

    return X

def forward(
    self,
    X: Float[torch.Tensor, "n_nodes dim"],
    mask: Float[torch.Tensor, "n_nodes n_nodes"] | None = None,
) -> Float[torch.Tensor, "n_nodes dim"]:
    """Forward pass for the stereographic attention layer."""
    Q = self._split_heads(self.W_q(X))  # [B, H, N, D]
    K = self._split_heads(self.W_k(X))
    V = self._split_heads(self.W_v(X=X))

    attn_out = self.attn(Q, K, V, mask.unsqueeze(0).unsqueeze(0))  # type: ignore
    attn_out = self._combine_heads(attn_out)

    out = self.ff(X=attn_out)

    return out

def forward(
    self,
    X: Float[torch.Tensor, "n_nodes dim"],
    mask: Float[torch.Tensor, "n_nodes n_nodes"] | None = None,
) -> Float[torch.Tensor, "n_nodes dim"]:
    """Forward pass through the stereographic transformer block."""
    X = geoopt.manifolds.stereographic.math.mobius_add(self.mha(self.norm1(X), mask), X, self.curvatures)
    X = geoopt.manifolds.stereographic.math.project(X, self.curvatures)
    X = geoopt.manifolds.stereographic.math.mobius_add(self.mlpblock(self.norm2(X)), X, self.curvatures)
    X = geoopt.manifolds.stereographic.math.project(X, self.curvatures)

    return X