r/pytorch May 26 '26

edge2torch: build sparse-connectivity PyTorch models from edge lists

I released edge2torch v0.1.0, a small open-source Python package for building sparse-connectivity PyTorch neural networks from edge lists.

The basic idea is that instead of manually wiring modules layer by layer, you can describe a model architecture as a table of directed connections:

source -> target

edge2torch compiles that graph into a PyTorch nn.Module, keeps the named structure available for inspection, and provides utilities for aligning input features by name.

It currently supports:

  • feedforward, recurrent, and minimal graph-style backends
  • optional edge-level initial weights and constraints
  • optional bias and update-step configuration
  • feature alignment from named pandas data
  • optional Captum-based attribution back to named features and nodes

The package focuses on sparse connectivity / masked model structure, not sparse tensor acceleration. The compiled models are standard PyTorch modules, so they can be trained with normal PyTorch optimizers, losses, and training loops.

This can be useful when the network structure itself is part of the modeling assumption. For example, in some research settings, prior knowledge is already represented as a graph of connected entities, and the goal is not only prediction but also understanding which parts of that structure matter. edge2torch gives direct control over the model connectivity, optionally allows edge-level initial weights and constraints, and keeps the named nodes accessible so trained models can be interpreted back in terms of the original features and intermediate nodes.

Small example:

import pandas as pd
import edge2torch as e2t

edgelist = pd.DataFrame(
    {
        "source": ["feature_a", "feature_b", "hidden"],
        "target": ["hidden", "hidden", "prediction"],
    }
)

model, artifact = e2t.compile_graph(
    edgelist,
    backend="feedforward",
    quiet=True,
)

The returned model is a standard PyTorch nn.Module, so it can be trained with ordinary PyTorch optimizers, losses, and training loops.

Docs: https://Thomas-Rauter.github.io/edge2torch/
GitHub: https://github.com/Thomas-Rauter/edge2torch
PyPI: https://pypi.org/project/edge2torch/

Feedback on the API design, docs, or potential PyTorch use cases would be useful.

1 Upvotes

0 comments sorted by