holonomy-lib 0.8.0__py3-none-any.whl

holonomy_lib/optimization/__init__.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: FSL-1.1-ALv2
+# Copyright 2026 Synoros
+
+"""holonomy_lib.optimization: Riemannian optimizers on the library's manifolds.
+
+Currently implemented:
+  RiemannianSGD(manifold, lr)  : steepest descent on a manifold.
+
+The optimizer only calls `egrad_to_rgrad` and `retraction` on its manifold,
+so it works with any manifold that has both (`FixedRankManifold`,
+`SPDManifold`, `ProductManifold`). The flow is:
+
+  1. Caller computes the ambient gradient `g` of the objective.
+  2. `optimizer.step(point, g)` converts `g` using the metric at
+     `point`, scales by `-lr`, and retracts back onto the manifold.
+  3. Returns the new point. Caller drives the loop.
+
+The functional helper `riemannian_sgd_step` exposes the same logic
+without state, useful for tests and for integrating into custom
+training loops.
+
+Why no Adam (or RMSProp, AdamW, ...): adaptive step-size schemes are
+user-side ergonomics, not part of the math of optimization on a
+manifold. SGD uses metric gradient conversion followed by retraction.
+Users who want Adam-style rates can compute their adaptive step in
+ambient space and pass the rescaled gradient into `step()`.
+
+Planned:
+  - `torch.optim.Optimizer` integration via a `ManifoldParameter`
+    wrapper so manifold params can be mixed with ambient params in a
+    normal PyTorch training loop. Out of scope for v1.
+  - Trust-region and conjugate-gradient solvers.
+"""
+
+from holonomy_lib.optimization.base import RiemannianOptimizer
+from holonomy_lib.optimization.sgd import RiemannianSGD, riemannian_sgd_step
+
+__all__ = [
+    "RiemannianOptimizer",
+    "RiemannianSGD",
+    "riemannian_sgd_step",
+]

holonomy_lib/optimization/base.py

@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: FSL-1.1-ALv2
+# Copyright 2026 Synoros
+
+"""Base class for Riemannian optimizers.
+
+Subclasses implement `step(point, ambient_grad) -> new_point`. The base
+class holds only the manifold reference. Hyperparameters and any per-step
+state belong to the subclass.
+
+The "manifold" parameter must expose `egrad_to_rgrad(point, ambient)` and
+`retraction(point, tangent)` methods with the same signatures as
+`FixedRankManifold` and `SPDManifold`. The optimizer is otherwise
+agnostic to which manifold it's operating on.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class RiemannianOptimizer(ABC):
+    """Abstract base for Riemannian optimizers.
+
+    An optimizer instance holds:
+      - the manifold (stored by `__init__`; provides gradient conversion and
+        retraction),
+      - hyperparameters (learning rate, momentum coefficients, etc.),
+      - per-instance state (moment buffers, step count) if stateful.
+
+    The subclass implements `step(point, ambient_grad)`, returning the
+    new point on the manifold. Caller is responsible for storing the
+    new point and reading its gradient on the next iteration.
+
+    Why this shape: the `FixedRankManifold` and `SPDManifold` methods
+    return new points rather than mutating in place, so the optimizer
+    follows the same convention. Wrapping a `torch.optim.Optimizer` around
+    this, for `loss.backward()` flows, is possible later. The explicit
+    point-in / point-out contract is the simpler primitive to build on.
+    """
+
+    def __init__(self, manifold: Any) -> None:
+        self.manifold = manifold
+
+    @abstractmethod
+    def step(self, point: Any, ambient_grad: Any) -> Any:
+        """Apply one optimizer step and return the new point.
+
+        Args:
+          point: current point on the manifold (tensor for SPD, triple
+            of tensors for FixedRank).
+          ambient_grad: gradient of the objective with respect to the
+            ambient (embedding-space) coordinates. Same structure as
+            an ambient tangent: a single tensor for both SPD and
+            FixedRank (FixedRank tangents are stored in ambient
+            `(B, m, n)` form per Vandereycken 2013).
+
+        Returns:
+          new_point: the point after the step.
+
+        References:
+          Bonnabel (2013), Stochastic Gradient Descent on Riemannian
+          Manifolds, §III.
+        """
+        ...

holonomy_lib/optimization/sgd.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: FSL-1.1-ALv2
+# Copyright 2026 Synoros
+
+"""Riemannian steepest descent (RSGD).
+
+For a Riemannian manifold M with `egrad_to_rgrad` (ambient gradient to
+Riemannian gradient, metric included) and `retraction` (move along a
+tangent vector to a new point), Riemannian SGD takes the ambient
+gradient `g` at the current point, converts it with the metric, scales
+by `-lr`, and retracts back onto M.
+
+  ξ = egrad_to_rgrad(point, g)
+  point' = retraction(point, -lr · ξ)
+
+References:
+  Absil, P.-A., Mahony, R., Sepulchre, R. (2008). Optimization
+    Algorithms on Matrix Manifolds. Princeton University Press,
+    §4.1-§4.2.
+  Bonnabel, S. (2013). Stochastic gradient descent on Riemannian
+    manifolds. IEEE TAC 58(9):2217-2229.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from holonomy_lib.optimization.base import RiemannianOptimizer
+
+
+def riemannian_sgd_step(
+    manifold: Any,
+    point: Any,
+    ambient_grad: torch.Tensor,
+    lr: float,
+) -> Any:
+    """Functional one-step Riemannian SGD.
+
+    Args:
+      manifold: object exposing `egrad_to_rgrad(point, ambient)` and
+        `retraction(point, tangent)` methods (e.g.
+        `FixedRankManifold`, `SPDManifold`).
+      point: current point on the manifold (tensor for SPD; triple
+        of tensors for FixedRank).
+      ambient_grad: ambient-space gradient of the objective. Always
+        a single `torch.Tensor`; for `FixedRankManifold` this is the
+        `(B, m, n)` ambient gradient per Vandereycken (2013)'s
+        embedded-tangent representation.
+      lr: learning rate, a positive scalar. Not checked here;
+        `RiemannianSGD` rejects lr <= 0.
+
+    Returns:
+      new_point: the manifold-constrained next iterate.
+
+    References:
+      Absil, P.-A., Mahony, R., Sepulchre, R. (2008). Optimization
+        Algorithms on Matrix Manifolds. Princeton University Press,
+        §4.1; retraction-based steepest descent.
+      Bonnabel, S. (2013). Stochastic gradient descent on Riemannian
+        manifolds. IEEE TAC 58(9):2217-2229, §III.
+    """
+    # A plain tangent projection is the right conversion only when the metric
+    # is the embedded Euclidean one, as for FixedRankManifold. SPD's
+    # affine-invariant metric needs S sym(G) S, so each manifold does its own
+    # conversion through egrad_to_rgrad. The step then scales that tangent by
+    # -lr and retracts it back onto the manifold.
+    tangent = manifold.egrad_to_rgrad(point, ambient_grad)
+    return manifold.retraction(point, -lr * tangent)
+
+
+class RiemannianSGD(RiemannianOptimizer):
+    """Optimizer that holds a manifold and a learning rate.
+
+    Each `step` call forwards to `riemannian_sgd_step`. The object keeps no
+    moment buffers or step count, so nothing carries over between steps.
+
+    Args:
+      manifold: the underlying manifold object (see `riemannian_sgd_step`).
+      lr: learning rate, must be > 0 (ValueError otherwise). The default
+        `1e-2` is a starting point, not a tuned value; tune it per problem.
+
+    Example:
+      >>> from holonomy_lib.manifolds import SPDManifold
+      >>> mfd = SPDManifold(n=4)
+      >>> opt = RiemannianSGD(mfd, lr=0.05)
+      >>> point = mfd.random_point(batch_size=1)
+      >>> # in a training loop:
+      >>> ambient_grad = compute_ambient_grad(point)
+      >>> point = opt.step(point, ambient_grad)
+    """
+
+    def __init__(self, manifold: Any, lr: float = 1e-2) -> None:
+        super().__init__(manifold)
+        if lr <= 0:
+            raise ValueError(f"lr must be > 0, got {lr}")
+        self.lr = lr
+
+    def step(self, point: Any, ambient_grad: Any) -> Any:
+        """Convert the ambient gradient and retract one descent step.
+
+        Args:
+          point: current point on the manifold.
+          ambient_grad: ambient-space gradient of the objective at `point`.
+        Returns:
+          new_point: the point after the step.
+
+        References:
+          Bonnabel (2013), Stochastic Gradient Descent on Riemannian
+          Manifolds, §III.
+        """
+        return riemannian_sgd_step(self.manifold, point, ambient_grad, self.lr)

holonomy_lib/sheaf/__init__.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: FSL-1.1-ALv2
+# Copyright 2026 Synoros
+
+"""holonomy_lib.sheaf; cellular sheaves and sheaf Laplacians on graphs.
+
+A **cellular sheaf** on a graph attaches a finite-dimensional vector
+space (a "stalk") `F(σ)` to each simplex σ (node or edge) and a linear
+**restriction map** `F_{τ ≤ σ}: F(τ) → F(σ)` for each face relation
+τ ≤ σ. The sheaf is the data of those stalks + maps; the sheaf
+Laplacian generalizes the graph Laplacian by capturing the
+"disagreement" of node-stalk values once they are pushed up to each
+incident edge stalk:
+
+    L_F = δ^T δ,   δ_e(x)  =  F_{u ≤ e}(x_u) − F_{v ≤ e}(x_v)
+                              for each edge e = (u, v).
+
+When all stalks have dimension 1 and every restriction map is the
+identity (i.e., the *trivial sheaf*), `L_F` reduces to the standard
+combinatorial graph Laplacian. With higher-dimensional stalks and
+non-trivial restriction maps it captures structure that a scalar
+Laplacian cannot; orientation sheaves, signed cellular sheaves,
+connection Laplacians, and the Laplacians used by Neural Sheaf
+Diffusion (Bodnar et al. 2022).
+
+Currently exposed:
+  GraphSheaf                 ; dataclass holding edges + stalks + maps
+  HeterogeneousGraphSheaf    ; per-node (heterogeneous) stalk dimensions
+  sheaf_coboundary(sheaf)    ; δ as a (n_e·d_e, n_v·d_v) tensor
+  sheaf_laplacian(sheaf)     ; L_F = δ^T δ, (n_v·d_v, n_v·d_v), PSD
+  sheaf_coboundary_sparse(sheaf); sparse COO coboundary
+  sheaf_laplacian_sparse(sheaf); sparse COO Laplacian
+  sheaf_bottom_spectrum(sheaf, k); bottom-k spectrum with optional kernel removal
+  sheaf_dirichlet_energy(sheaf, x)
+                             ; `x^T L_F x`, batched-first
+
+Restricted to **node-edge sheaves on graphs**. Higher-dimensional
+cellular sheaves on simplicial complexes (with 2-cells / faces and
+the corresponding ∂_1 ∘ ∂_2 = 0 chain identity) are not implemented.
+
+References:
+  Hansen, J., Ghrist, R. (2019). Toward a spectral theory of cellular
+    sheaves. Journal of Applied and Computational Topology 3:315–358.
+    The canonical reference for sheaf Laplacians and their spectral
+    theory.
+  Bodnar, C., Di Giovanni, F., Chamberlain, B. P., Liò, P., Bronstein,
+    M. M. (2022). Neural sheaf diffusion: A topological perspective on
+    heterophily and oversmoothing in GNNs. NeurIPS 2022. Brought
+    sheaf Laplacians into mainstream geometric deep learning.
+  Curry, J. (2014). Sheaves, Cosheaves and Applications. PhD thesis,
+    University of Pennsylvania. The applied-topology textbook on
+    cellular sheaves; introduces the explicit cellular formulation
+    that all the computational work uses.
+"""
+
+from holonomy_lib.sheaf.graph_sheaf import GraphSheaf, HeterogeneousGraphSheaf
+from holonomy_lib.sheaf.laplacian import (
+    sheaf_bottom_spectrum,
+    sheaf_coboundary,
+    sheaf_coboundary_sparse,
+    sheaf_dirichlet_energy,
+    sheaf_laplacian,
+    sheaf_laplacian_sparse,
+)
+
+__all__ = [
+    "GraphSheaf",
+    "HeterogeneousGraphSheaf",
+    "sheaf_bottom_spectrum",
+    "sheaf_coboundary",
+    "sheaf_coboundary_sparse",
+    "sheaf_dirichlet_energy",
+    "sheaf_laplacian",
+    "sheaf_laplacian_sparse",
+]

holonomy_lib/sheaf/graph_sheaf.py

@@ -0,0 +1,352 @@
+# SPDX-License-Identifier: FSL-1.1-ALv2
+# Copyright 2026 Synoros
+
+"""GraphSheaf dataclass; the data of a cellular sheaf on a graph.
+
+The sheaf attaches a `d_v`-dim stalk to each node and a `d_e`-dim
+stalk to each edge. For every edge `e = (u, v)` it carries two
+restriction maps:
+
+  F_left[e]  : R^{d_v} → R^{d_e}   (the map from node u's stalk)
+  F_right[e] : R^{d_v} → R^{d_e}   (the map from node v's stalk)
+
+The coboundary on a node-stalk vector `x ∈ R^{n_v · d_v}` reads each
+edge's contribution as `F_left[e] · x_u − F_right[e] · x_v`. The
+sheaf Laplacian is `δ^T δ`.
+
+Representations:
+  - `GraphSheaf` stores uniform stalk dimensions in dense tensors.
+    `HeterogeneousGraphSheaf` stores per-cell dimensions and ragged maps.
+  - Undirected graphs: edges are unordered, but the (u, v) order in each
+    tuple sets the sign convention of `δ`. Flipping an edge to (v, u)
+    and swapping its left and right maps gives the same sheaf, with that
+    edge's block of `δ` negated. The Laplacian `L_F = δ^T δ` does not
+    change, since it is quadratic in `δ`.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass(frozen=True)
+class GraphSheaf:
+    """A cellular sheaf on a graph.
+
+    Attributes:
+      n_nodes:      number of nodes (each carries a d_v-dim stalk).
+      edges:        `(n_edges, 2)` int64 tensor of `(u, v)` pairs.
+      node_stalk_dim:  `d_v`.
+      edge_stalk_dim:  `d_e`.
+      F_left:       `(n_edges, d_e, d_v)` restriction maps for the
+                    `u`-side of each edge.
+      F_right:      `(n_edges, d_e, d_v)` restriction maps for the
+                    `v`-side of each edge.
+
+    Validation:
+      - `edges` must be `(n_edges, 2)` int64.
+      - `F_left`, `F_right` shapes match `(n_edges, d_e, d_v)`.
+      - All tensors share `device`; restriction maps share a real float dtype.
+
+    References:
+      Hansen, J., Ghrist, R. (2019). Toward a spectral theory of cellular
+        sheaves. Journal of Applied and Computational Topology 3:315–358, §3.
+    """
+
+    n_nodes: int
+    edges: torch.Tensor
+    node_stalk_dim: int
+    edge_stalk_dim: int
+    F_left: torch.Tensor
+    F_right: torch.Tensor
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.n_nodes, int) or isinstance(self.n_nodes, bool):
+            raise ValueError("n_nodes must be an integer")
+        if self.n_nodes <= 0:
+            raise ValueError(f"n_nodes must be > 0; got {self.n_nodes}")
+        if (not isinstance(self.node_stalk_dim, int)
+                or isinstance(self.node_stalk_dim, bool)):
+            raise ValueError("node_stalk_dim must be an integer")
+        if (not isinstance(self.edge_stalk_dim, int)
+                or isinstance(self.edge_stalk_dim, bool)):
+            raise ValueError("edge_stalk_dim must be an integer")
+        if self.node_stalk_dim <= 0:
+            raise ValueError(
+                f"node_stalk_dim must be > 0; got {self.node_stalk_dim}"
+            )
+        if self.edge_stalk_dim <= 0:
+            raise ValueError(
+                f"edge_stalk_dim must be > 0; got {self.edge_stalk_dim}"
+            )
+        if self.edges.ndim != 2 or self.edges.shape[-1] != 2:
+            raise ValueError(
+                f"edges must be (n_edges, 2); got "
+                f"shape={tuple(self.edges.shape)}"
+            )
+        n_e = self.edges.shape[0]
+        d_e, d_v = self.edge_stalk_dim, self.node_stalk_dim
+        for name, t in (("F_left", self.F_left), ("F_right", self.F_right)):
+            if t.shape != (n_e, d_e, d_v):
+                raise ValueError(
+                    f"{name} must have shape (n_edges={n_e}, "
+                    f"edge_stalk_dim={d_e}, node_stalk_dim={d_v}); got "
+                    f"shape={tuple(t.shape)}"
+                )
+        if self.edges.dtype != torch.int64:
+            raise ValueError("edges must have dtype int64")
+        if not self.F_left.is_floating_point() or not self.F_right.is_floating_point():
+            raise ValueError("restriction maps must have a real floating dtype")
+        if self.edges.device != self.F_left.device:
+            raise ValueError("edges and restriction maps must share device")
+        if self.F_left.device != self.F_right.device:
+            raise ValueError(
+                f"F_left and F_right must be on the same device; "
+                f"got {self.F_left.device} vs {self.F_right.device}"
+            )
+        if self.F_left.dtype != self.F_right.dtype:
+            raise ValueError(
+                f"F_left and F_right must share dtype; got "
+                f"{self.F_left.dtype} vs {self.F_right.dtype}"
+            )
+        # Negative indices would wrap around silently in torch indexing, so
+        # both ends of the range are checked here.
+        if (self.edges < 0).any() or (self.edges >= self.n_nodes).any():
+            raise ValueError(
+                f"edges entries must be in [0, n_nodes={self.n_nodes})"
+            )
+        # A self-loop adds a diagonal term (F_left - F_right)^T (F_left - F_right)
+        # to the sheaf Laplacian. The library drops self-loops from graphs
+        # (CONVENTIONS.md, via `_graph_utils.drop_self_loops`), so the sheaf
+        # code refuses them too. To proceed, drop them first:
+        # `mask = u != v; edges = edges[mask]; F_left = F_left[mask];
+        # F_right = F_right[mask]`.
+        if (self.edges[:, 0] == self.edges[:, 1]).any():
+            raise ValueError(
+                "self-loops (u == v in edges) are not supported in v1; "
+                "drop them before constructing the GraphSheaf"
+            )
+        # Duplicate unordered endpoint pairs would silently double the
+        # off-diagonal mass of the sheaf Laplacian, breaking the
+        # "trivial sheaf → graph Laplacian" reduction. Reject them.
+        unique_edges = torch.unique(self.edges.sort(dim=-1).values, dim=0)
+        if unique_edges.shape[0] != self.edges.shape[0]:
+            raise ValueError(
+                "duplicate edges are not supported in v1; "
+                "the resulting sheaf Laplacian would double-count "
+                "the duplicated pair's off-diagonal block"
+            )
+
+    @property
+    def n_edges(self) -> int:
+        return int(self.edges.shape[0])
+
+    @property
+    def total_node_dim(self) -> int:
+        """Length of the stacked node-stalk vector, `n_nodes · d_v`."""
+        return self.n_nodes * self.node_stalk_dim
+
+    @property
+    def device(self) -> torch.device:
+        return self.F_left.device
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.F_left.dtype
+
+
+    @staticmethod
+    def trivial(
+        n_nodes: int,
+        edges: torch.Tensor,
+        stalk_dim: int = 1,
+        device: str | torch.device | None = None,
+        dtype: torch.dtype = torch.float64,
+    ) -> GraphSheaf:
+        """Build the trivial sheaf: identity restriction maps on every edge.
+
+        With `stalk_dim = 1` this is the sheaf whose Laplacian equals the
+        standard combinatorial graph Laplacian. With `stalk_dim = d > 1`,
+        node-major ordering gives `L ⊗ I_d`, which is permutation-similar to
+        a block diagonal matrix of `d` copies of L.
+
+        Args:
+          n_nodes: number of nodes.
+          edges: `(n_edges, 2)` int64 `(u, v)` pairs.
+          stalk_dim: stalk dimension `d` on every node and edge; must be > 0.
+          device: target device. Defaults to the device of `edges`, and an
+            explicit value moves both the indices and the maps there.
+          dtype: float dtype of the restriction maps.
+
+        Returns:
+          A `GraphSheaf` with `node_stalk_dim == edge_stalk_dim == stalk_dim`.
+
+        References:
+          Hansen-Ghrist (2019), §3.1, constant sheaves and graph Laplacians.
+        """
+        if edges.ndim != 2 or edges.shape[-1] != 2:
+            raise ValueError(
+                f"edges must be (n_edges, 2); got shape={tuple(edges.shape)}"
+            )
+        if edges.dtype != torch.int64:
+            raise ValueError("edges must have dtype int64")
+        if not isinstance(stalk_dim, int) or isinstance(stalk_dim, bool):
+            raise ValueError("stalk_dim must be an integer")
+        if stalk_dim <= 0:
+            raise ValueError("stalk_dim must be > 0")
+        if not dtype.is_floating_point:
+            raise ValueError("restriction maps must have a real floating dtype")
+        n_e = edges.shape[0]
+        device = edges.device if device is None else device
+        eye = torch.eye(stalk_dim, dtype=dtype, device=device)
+        F = eye.unsqueeze(dim=0).expand(n_e, stalk_dim, stalk_dim).contiguous()
+        edges = edges.to(device=device, dtype=torch.int64)
+        # F_left and F_right get separate copies, so writing into one of them
+        # does not change the other.
+        return GraphSheaf(
+            n_nodes=n_nodes,
+            edges=edges,
+            node_stalk_dim=stalk_dim,
+            edge_stalk_dim=stalk_dim,
+            F_left=F.clone(),
+            F_right=F.clone(),
+        )
+
+
+@dataclass(frozen=True)
+class HeterogeneousGraphSheaf:
+    """A cellular sheaf with PER-NODE (heterogeneous) stalk dimensions.
+
+    The generalization of `GraphSheaf` to per-node stalk dimensions: node `i`
+    carries a `d_v[i]`-dim stalk and edge `e` a `d_e[e]`-dim stalk, all free
+    to differ. For edge `e = (u, v)`:
+
+      F_left[e]  : R^{d_v[u]} -> R^{d_e[e]}   shape (d_e[e], d_v[u])
+      F_right[e] : R^{d_v[v]} -> R^{d_e[e]}   shape (d_e[e], d_v[v])
+
+    Coboundary, Laplacian, and Dirichlet energy mean exactly what they do in the
+    uniform case (each edge block of `δ` is `F_left[e]·x_u − F_right[e]·x_v`,
+    `L_F = δ^T δ`); only the per-node/per-edge dims and the ragged restriction
+    maps differ. The maps are stored as a tuple of tensors (ragged), since the
+    dims vary, rather than one dense `(n_e, d_e, d_v)` block.
+
+    Use this when a node's stalk dimension is a quantity of its own, such as
+    each concept's relational rank, rather than one global K.
+
+    Attributes:
+      n_nodes:         number of nodes.
+      edges:           `(n_edges, 2)` int64 `(u, v)` pairs.
+      node_stalk_dims: `(n_nodes,)` int64, per-node stalk dim `d_v[i] > 0`.
+      edge_stalk_dims: `(n_edges,)` int64, per-edge stalk dim `d_e[e] > 0`.
+      F_left:          length-`n_edges` tuple; `F_left[e]` is `(d_e[e], d_v[u])`.
+      F_right:         length-`n_edges` tuple; `F_right[e]` is `(d_e[e], d_v[v])`.
+
+    References:
+      Hansen-Ghrist (2019): cellular sheaves carry per-cell stalks of arbitrary
+        finite dimension; the uniform `GraphSheaf` is the special case.
+    """
+
+    n_nodes: int
+    edges: torch.Tensor
+    node_stalk_dims: torch.Tensor
+    edge_stalk_dims: torch.Tensor
+    F_left: tuple[torch.Tensor, ...]
+    F_right: tuple[torch.Tensor, ...]
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.n_nodes, int) or isinstance(self.n_nodes, bool):
+            raise ValueError("n_nodes must be an integer")
+        if self.n_nodes <= 0:
+            raise ValueError(f"n_nodes must be > 0; got {self.n_nodes}")
+        if self.edges.ndim != 2 or self.edges.shape[-1] != 2:
+            raise ValueError(
+                f"edges must be (n_edges, 2); got shape={tuple(self.edges.shape)}"
+            )
+        n_e = int(self.edges.shape[0])
+        if tuple(self.node_stalk_dims.shape) != (self.n_nodes,):
+            raise ValueError(
+                f"node_stalk_dims must be (n_nodes={self.n_nodes},); got "
+                f"shape={tuple(self.node_stalk_dims.shape)}"
+            )
+        if tuple(self.edge_stalk_dims.shape) != (n_e,):
+            raise ValueError(
+                f"edge_stalk_dims must be (n_edges={n_e},); got "
+                f"shape={tuple(self.edge_stalk_dims.shape)}"
+            )
+        for name, tensor in (("edges", self.edges),
+                             ("node_stalk_dims", self.node_stalk_dims),
+                             ("edge_stalk_dims", self.edge_stalk_dims)):
+            if tensor.dtype != torch.int64:
+                raise ValueError(f"{name} must have dtype int64")
+        if (self.node_stalk_dims <= 0).any() or (self.edge_stalk_dims <= 0).any():
+            raise ValueError("all stalk dims must be > 0")
+        if len(self.F_left) != n_e or len(self.F_right) != n_e:
+            raise ValueError(
+                f"F_left and F_right must each have n_edges={n_e} maps; got "
+                f"{len(self.F_left)} and {len(self.F_right)}"
+            )
+        if (self.edges < 0).any() or (self.edges >= self.n_nodes).any():
+            raise ValueError(
+                f"edges entries must be in [0, n_nodes={self.n_nodes})"
+            )
+        if (self.edges[:, 0] == self.edges[:, 1]).any():
+            raise ValueError(
+                "self-loops (u == v in edges) are not supported; drop them first"
+            )
+        if torch.unique(self.edges.sort(dim=-1).values, dim=0).shape[0] != n_e:
+            raise ValueError("duplicate edges are not supported")
+        # Each F_left[e] must be (d_e[e], d_v[u]) and each F_right[e] must be
+        # (d_e[e], d_v[v]). The shapes change from edge to edge, so this is
+        # checked one edge at a time. The sizes and endpoints are pulled into
+        # Python lists once, so the loop never reads the device tensors.
+        node_dims = self.node_stalk_dims.tolist()
+        edge_dims = self.edge_stalk_dims.tolist()
+        uv = self.edges.tolist()
+        for e in range(n_e):
+            u, v = uv[e]
+            want_left = (edge_dims[e], node_dims[u])
+            want_right = (edge_dims[e], node_dims[v])
+            if tuple(self.F_left[e].shape) != want_left:
+                raise ValueError(
+                    f"F_left[{e}] must be (d_e[e], d_v[u])={want_left}; got "
+                    f"{tuple(self.F_left[e].shape)}"
+                )
+            if tuple(self.F_right[e].shape) != want_right:
+                raise ValueError(
+                    f"F_right[{e}] must be (d_e[e], d_v[v])={want_right}; got "
+                    f"{tuple(self.F_right[e].shape)}"
+                )
+        # With no edges there are no maps to compare, so the device and dtype
+        # check only runs when at least one edge exists.
+        if n_e > 0:
+            dev, dt = self.F_left[0].device, self.F_left[0].dtype
+            for t in (*self.F_left, *self.F_right):
+                if not t.is_floating_point():
+                    raise ValueError("restriction maps must have a real floating dtype")
+                if t.device != dev or t.dtype != dt:
+                    raise ValueError(
+                        "all restriction maps must share device and dtype"
+                    )
+
+    @property
+    def n_edges(self) -> int:
+        return int(self.edges.shape[0])
+
+    @property
+    def total_node_dim(self) -> int:
+        """Length of the stacked node-stalk vector, `Σ_i d_v[i]`."""
+        return int(self.node_stalk_dims.sum())
+
+    # With no edges there is no map to read the device or dtype from. The
+    # device falls back to the stalk-dim tensor, and the dtype to float64,
+    # the default that `trivial` uses.
+    @property
+    def device(self) -> torch.device:
+        return (self.F_left[0].device if self.n_edges
+                else self.node_stalk_dims.device)
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.F_left[0].dtype if self.n_edges else torch.float64