Skip to content

edge2torch

ci codecov PyPI Python PyPI - License PyPI - Downloads pypi since

Build sparsely connected PyTorch neural networks from prior-knowledge graphs, with optional feature- and node-level attribution.

Graphical abstract of edge2torch

Overview

edge2torch is an edge-list-to-PyTorch compiler for sparsely-connected neural network architectures with named nodes.

An edge list is a table of directed connections: each row links a source node to a target node. For example:

source target
feature_a hidden_1
feature_b hidden_1
hidden_1 output

Define a model architecture as an edge list, compile it into a minimally opinionated PyTorch model, train it with standard PyTorch tools, and optionally map model behavior back to the named nodes and features that defined the architecture.

The package is designed for users who want to build sparsely-connected or structured neural networks from a predefined graph rather than manually wiring PyTorch modules. It is domain-agnostic: any setting where a neural architecture can be represented as named edges can use the same graph-to-model abstraction.

Here, "graph" means the architecture specification, not necessarily a graph neural network. Feedforward models and topology-preserving state-update models can both be represented by edge lists when their architecture is defined through directed connections between named nodes.

A major application area is interpretable neural networks shaped by prior knowledge of domain networks, for example in biology and chemistry (more about this on the Getting started page).

edge2torch deliberately leaves training loops, losses, optimizers, task-specific heads, advanced customization, etc. to standard PyTorch. This makes the package very flexible.

Core workflow

The package is built around four main steps:

  1. Define a model architecture as an edge list with named source and target nodes.
  2. Compile the edge list into a backend-specific PyTorch model with compile_graph().
  3. Align named input data features to the compiled model input nodes with align_features_to_input_nodes().
  4. Customize and train the model with ordinary PyTorch, or customize with customize_model().
  5. Interpret the model with attribution algorithms implemented in the Captum Python package, or use interpret_model() as a convenient interface for Captum.

Main public API

The current public API is centered on:

  • compile_graph()
  • align_features_to_input_nodes()
  • customize_model()
  • edge_weights()
  • interpret_model()

CompileArtifact is also exported as a public type because it is returned by compile_graph() and consumed by helper functions. Its stable public fields are backend, input_nodes, output_nodes, hidden_nodes, and interpretation_sites. feature_names is a read-only alias of input_nodes (same names, same order).

Optional helpers and types: InterpretationResult from interpret_model(), plus CompileBackend, COMPILE_BACKENDS, Edge2TorchError, and __version__.

See the API reference page for details.

Package philosophy

edge2torch is intentionally minimally opinionated.

It defines the structural semantics required to compile a graph into a neural network backend, but it does not impose broader modeling choices such as:

  • activation functions
  • output heads
  • dropout
  • loss functions
  • optimizers
  • training loops
  • etc.

These remain part of the normal PyTorch workflow:

  • edge2torch handles graph compilation and provides convenience functions
  • PyTorch handles model training
  • edge2torch maps trained models back to interpretable named entities

Supported backends

compile_graph() currently supports:

  • feedforward
  • state_update

These backends share the same edge-list input format but differ in how the graph structure is translated into neural-network computation.

Feature attribution and (hidden) node-level interpretations are available through Captum-based methods on both backends.

See the Backends page for details.

Start here

If you are new to the package, start with:

License

This project is licensed under the MIT License. See the LICENSE file on GitHub for details.