Layers¶

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)

    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."""
    return self.manifold.manifold.projx(self.norm(X))

def forward(
    self,
    Q: Float[torch.Tensor, "num_heads n_nodes head_dim"],
    K: Float[torch.Tensor, "num_heads n_nodes head_dim"],
    V: Float[torch.Tensor, "num_heads n_nodes head_dim"],
    curvatures: Float[torch.Tensor, "num_heads 1 1"],
    mask: Float[torch.Tensor, "n_nodes n_nodes"] | None = None,
) -> Float[torch.Tensor, "num_heads n_nodes head_dim"]:
    """Forward pass for faithful gyrovector linear attention.

    Args:
        Q: Query tangent vectors at ``V``, shape ``[num_heads, n_nodes, head_dim]``.
        K: Key tangent vectors at ``V``, shape ``[num_heads, n_nodes, head_dim]``.
        V: Value points on the per-head manifold, shape ``[num_heads, n_nodes, head_dim]``.
        curvatures: Per-head curvatures, shape ``[num_heads, 1, 1]``.
        mask: Optional adjacency/attention mask, shape ``[n_nodes, n_nodes]``. Entry ``(i, j)``
            weights how much query ``i`` attends to key/value ``j``. ``None`` means full attention.

    Returns:
        Aggregated value points on the per-head manifold, shape ``[num_heads, n_nodes, head_dim]``.
    """
    k = curvatures
    math = geoopt.manifolds.stereographic.math

    # Eq 6: parallel-transport Q, K (tangent at V) to the origin, then feature map phi = elu + 1.
    q_tilde = math.parallel_transport0back(V, Q, k=k)  # [H, N, d]
    k_tilde = math.parallel_transport0back(V, K, k=k)  # [H, N, d]
    phi_q = nn.functional.elu(q_tilde) + 1.0  # phi(Q~)
    phi_k = nn.functional.elu(k_tilde) + 1.0  # phi(K~)

    # Conformal factor lambda^kappa at each value point (Eq 7).
    lam = math.lambda_x(x=V, k=k, keepdim=True, dim=-1)  # [H, N, 1]
    denom = geoopt.utils.clamp_abs(lam - 1, self._clamp_epsilon)  # (lambda - 1), [H, N, 1]

    v_tilde = (lam / denom) * V  # V~_i = lambda_i / (lambda_i - 1) * V_i
    phi_k_prime = denom * phi_k  # phi'(K~)_i = phi(K~)_i * (lambda_i - 1)

    if mask is None:
        # Kernelized O(N + M) form (Eq 11). The numerator effectively weights V_j by
        # alpha_ij * lambda_j = (phi(Q~)_i . phi'(K~)_j) * V~_j, and the normalizer by
        # alpha_ij * (lambda_j - 1) = phi(Q~)_i . phi'(K~)_j.
        context = torch.einsum("hnd,hne->hde", phi_k_prime, v_tilde)  # phi'(K)^T V~, [H, d, d]
        numerator = torch.einsum("hnd,hde->hne", phi_q, context)  # [H, N, d]
        normalizer = torch.einsum("hnd,hd->hn", phi_q, phi_k_prime.sum(dim=-2))  # [H, N]
    else:
        # Explicit (masked) attention scores -- supports a sparse adjacency, O(N^2 d).
        # Identical math to the kernelized path: numerator weights V_j by alpha_ij * lambda_j,
        # normalizer by alpha_ij * (lambda_j - 1).
        alpha = torch.einsum("hnd,hmd->hnm", phi_q, phi_k) * mask[None]  # [H, N, N]
        numerator = torch.einsum("hnm,hme->hne", alpha, lam * V)  # [H, N, d]
        normalizer = torch.einsum("hnm,hm->hn", alpha, denom.squeeze(-1))  # [H, N]

    norm_inv = 1.0 / normalizer.masked_fill(normalizer == 0, self._epsilon)  # [H, N]
    midpoint = numerator * norm_inv.unsqueeze(-1)  # weighted gyro-sum, [H, N, d]

    # Einstein midpoint finishes with a Mobius half-scaling; project for numerical safety.
    out = math.project(midpoint, k=k)
    out = math.mobius_scalar_mul(torch.tensor(0.5, dtype=out.dtype, device=out.device), out, k=k, dim=-1)
    out = math.project(out, k=k)
    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 for the mixed-curvature attention layer."""
    math = geoopt.manifolds.stereographic.math
    k = self._per_head_curvatures()  # [H, 1, 1]

    # Tangent space at the origin of the input manifold, then per-head Euclidean projections.
    H = self.manifold.manifold.logmap0(X)  # [N, dim]
    q_tan = self._split_heads(self.W_q(H))  # [H, N, d] (Euclidean, at origin)
    k_tan = self._split_heads(self.W_k(H))
    v_tan = self._split_heads(self.W_v(H))

    # Values are points on the per-head st_{kappa_h}; Q/K are tangent vectors AT each value point.
    V = math.expmap0(v_tan, k=k)  # [H, N, d] points on st_kappa_h
    Q = math.parallel_transport0(V, q_tan, k=k)  # tangent at V_i (Eq 5)
    K = math.parallel_transport0(V, k_tan, k=k)

    attn_out = self.attn(Q, K, V, curvatures=k, mask=mask)  # [H, N, d] points on st_kappa_h

    # Multi-head output is the product over heads: concat per-head tangent coords, project back.
    attn_tan = self._combine_heads(math.logmap0(attn_out, k=k))  # [N, inner] (tangent)
    out_tan = self.W_o(attn_tan)  # [N, dim]
    return self.manifold.manifold.expmap0(out_tan)

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 mixed-curvature transformer block.

    Args:
        X: Node features as points on the manifold, shape ``[n_nodes, dim]``.
        mask: Optional adjacency matrix ``A_hat``; ``None`` means full attention.

    Returns:
        Updated node features as points on the manifold, shape ``[n_nodes, dim]``.
    """
    man = self.manifold.manifold

    # Pre-norm attention sublayer with Mobius-addition residual.
    attn = self.mha(self.norm1(X), mask)
    X = man.projx(man.mobius_add(attn, X))

    # Pre-norm feedforward sublayer with Mobius-addition residual.
    ff = self._mlpblock(self.norm2(X))
    X = man.projx(man.mobius_add(ff, X))

    return X