API reference

Names exported from edge2torch. Each return object is documented immediately after the function that produces it.

compile_graph

Returns (model, artifact). model is a PyTorch nn.Module. artifact is a CompileArtifact.

Compile an edgelist into a sparse PyTorch model and compilation artifact.

The edgelist defines the architecture graph. Each row describes one directed connection from source to target in the direction of computation. In other words, edges should point from input feature nodes toward hidden nodes and output nodes.

Input features are inferred as graph nodes with no incoming edges. Output nodes are inferred as graph nodes with no outgoing edges. The returned artifact stores those input names in artifact.input_nodes and the read-only alias artifact.feature_names. Tensors passed to the compiled model must have columns in that exact order.

The edgelist may optionally include edge-level parameter metadata using the columns "initial_weight" and "constraint". These columns allow individual edges to define their initial effective weight and, where supported by the selected backend, constrain the trainable edge weight during optimization. Supported constraint values are "unconstrained", "positive", "negative", and "fixed". If omitted, edges use the backend's default trainable weight initialization and unconstrained weight behavior.

Graph-derived connectivity is enforced through masks on trainable edge weights. By default, compiled layers also include bias terms. Biases are node-level parameters, not graph edges, and are not constrained by the edge mask. Set bias=False to remove these offsets so node updates depend only on graph-defined weighted inputs.

The state_update backend applies a fixed number of graph state-update steps during each forward pass. The steps argument controls this update count and defaults to 3 when omitted. Passing steps with the feedforward backend raises an error. It is not a training epoch count and does not represent a sequence length in the input data.

Parameters:

Name	Type	Description	Default
edgelist	DataFrame	Edge table with required columns "source" and "target". Each row defines a directed connection from one named node to another, following the direction of computation. The table must include edges from input feature nodes into the rest of the architecture graph. The table may also include optional columns "initial_weight" and "constraint". If provided, "initial_weight" defines the initial effective edge weight, and "constraint" defines how that edge weight is parameterized during training. Supported constraints are: "unconstrained": the edge weight is trainable and may become positive or negative. "positive": the edge weight is trainable and constrained to remain positive. "negative": the edge weight is trainable and constrained to remain negative. "fixed": the edge weight is fixed to "initial_weight" and is not trainable. The optional columns are independent and row-wise sparse. Missing "initial_weight" values use the backend's default initialization for that edge. Missing "constraint" values are treated as "unconstrained" for that edge. If both values are provided for an edge, positive-constrained edges must have positive initial weights and negative-constrained edges must have negative initial weights. Edges with constraint="fixed" must provide an "initial_weight" value in the same row, because fixed edges require an explicit constant value.	required
backend	str	Backend to compile to. One of "feedforward" or "state_update". Typed as CompileBackend.	"feedforward"
bias	bool	Whether compiled masked linear layers include bias terms. If True, each target node has a learned node-level offset in addition to its graph-defined weighted inputs. If False, node updates are computed only from graph-defined weighted inputs. Disabling bias gives the graph structure stricter control over node activations.	True
steps	int | None	Number of state-update steps for the state_update backend. Larger values allow information to propagate through longer graph paths and revisit cycles more times. When omitted, state_update uses 3 steps. Must be omitted for the feedforward backend.	None

Returns:

Type	Description
tuple[Module, CompileArtifact]	Tuple (model, artifact). model is a PyTorch nn.Module compiled from the edgelist. artifact stores compilation metadata, including artifact.input_nodes and artifact.feature_names.

Raises:

Type	Description
Edge2TorchError	If input validation, graph validation, or backend compilation fails.

Notes

Informational compile notes are emitted through the edge2torch logger at INFO level. Graph warnings use warnings.warn(). Default logging configuration does not print INFO records, so notes are silent unless the caller configures logging.

Examples:

Compile a small feedforward architecture from an edgelist.

>>> import pandas as pd
>>> from edge2torch import compile_graph
>>>
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["feature_a", "feature_b", "hidden_1"],
...         "target": ["hidden_1", "hidden_1", "prediction"],
...     }
... )
>>>
>>> model, artifact = compile_graph(
...     edgelist=edgelist,
...     backend="feedforward",
... )
>>>
>>> artifact.feature_names
['feature_a', 'feature_b']

Compile a state-update architecture with edge-level initial weights and constraints.

>>> edgelist = pd.DataFrame(
...     {
...         "source": ["feature_a", "feature_b", "hidden_1"],
...         "target": ["hidden_1", "hidden_1", "prediction"],
...         "initial_weight": [0.1, -0.2, 0.5],
...         "constraint": ["positive", "negative", "fixed"],
...     }
... )
>>>
>>> model, artifact = compile_graph(
...     edgelist=edgelist,
...     backend="state_update",
... )

CompileArtifact

Returned by compile_graph(). Consumed by align_features_to_input_nodes() and interpret_model().

Public fields: backend, input_nodes, output_nodes, hidden_nodes, interpretation_sites. feature_names is a read-only alias of input_nodes.

Compilation metadata returned together with the compiled PyTorch model.

CompileArtifact is returned by compile_graph() and accepted by public helper functions such as align_features_to_input_nodes() and interpret_model(). It is exported for user-facing type hints and workflow integration.

The stable public fields are backend, input_nodes, output_nodes, hidden_nodes, and interpretation_sites. feature_names is a read-only alias of input_nodes: the same names in the same order, used as the tensor-column contract. interpretation_sites lists visible graph nodes only. Compilation internals such as the full attribution site map, the graph object, and the execution plan are private and may change across releases.

Parameters:

Name	Type	Description	Default
backend	CompileBackend	Backend used for compilation ("feedforward" or "state_update").	required
input_nodes	list[str]	Names of graph input nodes inferred as nodes with no incoming edges. This list is also the expected input column order for tensors passed to the compiled model.	required
output_nodes	list[str]	Names of graph output nodes inferred as nodes with no outgoing edges.	required
hidden_nodes	list[str]	Names of hidden graph nodes excluding inputs, outputs, and compiler pseudo nodes.	required
interpretation_sites	dict[str, list[str]]	Mapping from interpretation site identifier to ordered visible domain node names for that site. The feedforward backend uses layer_1, layer_2, and so on. The state_update backend uses step_1, step_2, and so on. Compiler pseudo nodes and other internal nodes are omitted.	required

align_features_to_input_nodes

Consumes the CompileArtifact from compile_graph(). Returns a torch.Tensor with columns in artifact.feature_names order.

Align data features to the input-node order expected by a compiled model.

compile_graph() builds a sparse neural network from an edgelist. Input nodes are inferred from the graph structure and stored in artifact.input_nodes. artifact.feature_names is a read-only alias of that list and defines the required column order for tensors passed to the compiled PyTorch model.

For named data containers, this function validates that every required input-node feature is present and reorders those features by name:

• pandas.DataFrame inputs are aligned using column names.
• AnnData inputs are aligned using var_names if anndata is installed.

Named data containers must contain the compiled model input-node features, although they may appear in any order. Missing features raise an error. Extra columns or AnnData variables are ignored. Non-string labels (for example integer DataFrame columns) are matched after converting to strings, matching how node IDs are stored after compile_graph().

torch.Tensor inputs do not contain feature names, so they are only validated by shape and are assumed to already follow artifact.feature_names order.

Parameters:

Name	Type	Description	Default
data	DataFrame | Tensor | AnnData	Input data to align. AnnData is supported when anndata is installed.	required
artifact	CompileArtifact	Compilation artifact returned by compile_graph(). artifact.feature_names is a read-only alias of artifact.input_nodes and defines the required input-node order.	required

Returns:

Type	Description
Tensor	Float32 input tensor whose columns are ordered according to artifact.feature_names.

Raises:

Type	Description
Edge2TorchError	If the input data type is unsupported, required features are missing, non-numeric DataFrame columns are present among the required features, or tensor input has an incompatible shape.

Examples:

Align a DataFrame whose columns are named but not ordered like the compiled model input nodes.

>>> import pandas as pd
>>> import torch
>>> from edge2torch import align_features_to_input_nodes, compile_graph
>>>
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["feature_a", "feature_b", "hidden"],
...         "target": ["hidden", "hidden", "prediction"],
...     }
... )
>>> model, artifact = compile_graph(edgelist)
>>>
>>> data = pd.DataFrame(
...     {
...         "feature_b": [2.0, 4.0],
...         "feature_a": [1.0, 3.0],
...     }
... )
>>>
>>> artifact.feature_names
['feature_a', 'feature_b']
>>>
>>> x = align_features_to_input_nodes(
...     data=data,
...     artifact=artifact,
... )
>>> x
tensor([[1., 2.],
        [3., 4.]])

Tensor inputs do not contain feature names, so they are only checked by shape and are assumed to already follow artifact.feature_names.

>>> x_tensor = torch.tensor(
...     [
...         [1.0, 2.0],
...         [3.0, 4.0],
...     ]
... )
>>> x_from_tensor = align_features_to_input_nodes(
...     data=x_tensor,
...     artifact=artifact,
... )
>>> torch.equal(x_from_tensor, x_tensor)
True

customize_model

Wraps the model returned by compile_graph(). Returns an nn.Module. The artifact is unchanged.

Wrap a compiled sparse neural network with optional PyTorch modules.

This function is a convenience layer for common post-compilation additions. It applies the requested components sequentially to the output of the compiled model. It does not modify the sparse graph structure, insert modules inside graph-derived layers, or replace ordinary PyTorch training and customization.

customize_model() wraps the provided model. Calling it repeatedly creates nested wrappers; it does not replace earlier customization modules. To change a customization, call customize_model() again on the original compiled model.

Parameters:

Name	Type	Description	Default
model	Module	PyTorch model returned by compile_graph().	required
activation	Module | None	Optional PyTorch activation module applied after the compiled model. This should be an instantiated module such as nn.ReLU().	None
dropout	float | int | None	Optional dropout probability applied after the activation. Must satisfy 0 <= dropout < 1.	None
head	Module | None	Optional PyTorch module applied after dropout. This should be an instantiated module such as nn.Linear(...).	None

Returns:

Type	Description
Module	Wrapped PyTorch model with the requested post-compilation modules.

Raises:

Type	Description
Edge2TorchError	If any input is invalid.

Examples:

Add an activation function after the compiled sparse neural network.

>>> import pandas as pd
>>> from torch import nn
>>> from edge2torch import compile_graph, customize_model
>>>
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["feature_a", "feature_b", "hidden"],
...         "target": ["hidden", "hidden", "prediction"],
...     }
... )
>>> model, artifact = compile_graph(edgelist)
>>>
>>> customized_model = customize_model(
...     model=model,
...     activation=nn.ReLU(),
... )

Add an activation, dropout, and task-specific prediction head.

>>> customized_model = customize_model(
...     model=model,
...     activation=nn.ReLU(),
...     dropout=0.2,
...     head=nn.Linear(1, 1),
... )

edge_weights

Reads effective named-edge weights from a compile_graph() model or a customize_model() wrapper. Returns a DataFrame with columns source, target, and weight.

Return the effective forward-pass weight of each original graph edge.

Compiled models store weights in backend-specific masked linear layers. Unconstrained layers, constrained layers, and feedforward skip-edge expansion all use different internal layouts. This helper maps those internals back to the original named source / target edges.

Weights are the values used during the forward pass, including softplus/fixed transforms when edge constraints are present. Pseudo-node routing edges used internally by the feedforward backend are omitted.

The compiled core is found the same way node interpretation finds a site provider: compile_graph() models, customize_model() wrappers, and manual wrappers that register the compiled model as a submodule all work.

Parameters:

Name	Type	Description	Default
model	Module	PyTorch model returned by compile_graph(), optionally wrapped by customize_model() or another nn.Module that keeps the compiled model as a registered submodule.	required

Returns:

Type	Description
DataFrame	Table with columns source, target, and weight. Rows follow the original compiled edgelist order. weight is the effective scalar used for that edge in the forward pass.

Raises:

Type	Description
Edge2TorchError	If no compiled EdgeModel / StateUpdateEdgeModel is found in model, or an original edge cannot be mapped onto the compiled weights.

Examples:

>>> import pandas as pd
>>> from edge2torch import compile_graph, edge_weights
>>>
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["feature_a", "feature_b", "hidden"],
...         "target": ["hidden", "hidden", "prediction"],
...     }
... )
>>> model, _ = compile_graph(edgelist)
>>>
>>> table = edge_weights(model)
>>> list(table.columns)
['source', 'target', 'weight']

interpret_model

Consumes a compiled (optionally customized) model and its CompileArtifact. Returns an InterpretationResult.

Interpret a model compiled by edge2torch using a Captum attribution method.

Parameters:

Name	Type	Description	Default
model	Any	PyTorch model returned by compile_graph(), optionally customized and trained by the user.	required
artifact	Any	Compilation artifact returned by compile_graph().	required
data	DataFrame | AnnData | Tensor	Input data used for attribution. Named containers are aligned to artifact.feature_names: required features may appear in any order, extra columns or AnnData variables are ignored, and missing required features raise an error. Tensor inputs are validated by shape only and are assumed to already follow artifact.feature_names order.	required
on	str	What to attribute onto. Use "features" to attribute predictions to input features. Use "nodes" to attribute predictions to named graph nodes. This is not Captum's output-index target argument; pass that through attribute_kwargs={"target": ...}.	"features"
method	str | None	Captum attribution method name. Method names follow Captum class names exactly and are case-sensitive, for example "IntegratedGradients", "Saliency", "DeepLift", "LayerConductance", or "LayerIntegratedGradients". When omitted, on="features" uses IntegratedGradients and on="nodes" uses LayerConductance. An explicit method must be compatible with on. If an unsupported method is provided, edge2torch raises an error listing the supported method names.	None
constructor_kwargs	dict[str, Any] | None	Optional keyword arguments passed directly to the constructor of the selected Captum attribution class. These arguments are passed through unchanged and are not interpreted, validated, or modified by edge2torch. Refer to the Captum documentation for the selected method to determine which constructor arguments are supported.	None
attribute_kwargs	dict[str, Any] | None	Optional keyword arguments passed directly to the selected Captum method's attribute() call. These arguments are passed through unchanged and are not interpreted, validated, or modified by edge2torch. Refer to the Captum documentation for the selected method to determine which attribution arguments are supported.	None

Returns:

Type

Description

InterpretationResult

Attribution result for the selected on value. If on="features", use result.table. That DataFrame has rows as examples and columns as features. Values are signed Captum scores. If on="nodes", use result.sites for one table per interpretation site. Each table contains the visible graph nodes the compiler exposes at that site (pseudo and other internal nodes are omitted). Values are signed Captum scores. Site keys are layer_* on feedforward and step_* on state_update. Feedforward sites start after the input layer, so input nodes do not appear in result.sites. State-update sites include all visible graph nodes, including inputs. Use result.summary() for one collapsed example-by-node table. summary() defaults to hidden nodes and aggregation="peak". Pass nodes="non_input" to include outputs, or nodes="all" to keep every visible node that appears in result.sites. nodes="all" includes inputs only on state_update. On feedforward, input attributions stay in result.table. Pass aggregation="mean_abs" or aggregation="last" to change how state-update steps are combined. "peak" and "last" keep the sign; "mean_abs" averages magnitudes. On feedforward results, a non-default aggregation is ignored and emits a warning.

Notes

Feature and node interpretation are supported for both backends (feedforward and state_update). Node interpretation uses Captum layer attribution classes at layer_* or step_* sites.

For node-level interpretation, edge2torch must access the compiled model's internal interpretation sites. This works for raw models returned by compile_graph(), models returned by customize_model(), and manually wrapped PyTorch models if the compiled model remains registered as a PyTorch submodule. Highly custom wrappers that hide, replace, or bypass the compiled model may not support on="nodes".

interpret_model() temporarily switches the model to evaluation mode while computing attributions and restores the previous training/evaluation mode afterward.

constructor_kwargs and attribute_kwargs are passed through to Captum. Refer to the Captum documentation for method-specific arguments such as baselines, targets, additional forward arguments, or perturbation settings.

Raises:

Type	Description
Edge2TorchError	If interpretation input validation fails, the requested on / method / backend combination is not supported, or Captum returns unsupported output.

Examples:

Compute feature-level attributions. The default method is integrated gradients.

>>> result = interpret_model(
...     model=trained_model,
...     artifact=artifact,
...     data=data,
...     on="features",
... )
>>> result.table.head()

Compute a summary of node-level attributions. The default method is layer conductance.

>>> result = interpret_model(
...     model=trained_model,
...     artifact=artifact,
...     data=data,
...     on="nodes",
... )
>>> result.summary().head()

Inspect per-site node-level attributions, including output nodes.

>>> result = interpret_model(
...     model=trained_model,
...     artifact=artifact,
...     data=data,
...     on="nodes",
...     method="LayerConductance",
... )
>>> result.sites.keys()
>>> result.summary(nodes="non_input").head()

InterpretationResult

Returned by interpret_model(). Wrong-on accessors raise Edge2TorchError.

• on="features": result.table (examples × features)
• on="nodes": result.sites and result.summary()
• sites: one table per site (layer_* on feedforward, step_* on state_update). Visible nodes at that site only. Feedforward sites start after the input layer (no inputs). State-update sites include inputs.
• summary(nodes=..., aggregation=...): one examples × nodes table. Defaults nodes="hidden", aggregation="peak". nodes="non_input" adds outputs. nodes="all" keeps every column in sites (inputs only on state_update).
• Scores in result.table and result.sites are signed Captum values. Nothing is abs'd at compute time. Rankings or plots that want magnitude should call .abs() themselves.
• aggregation is state_update only (peak / mean_abs / last). "peak" (default) and "last" keep the sign. "mean_abs" averages magnitudes and is the only option that drops the sign. Feedforward concatenates disjoint site columns; a non-default aggregation is ignored and warns.

Attribution tables produced by interpret_model().

Feature results expose table. Node results expose sites and summary(). Accessing the wrong on value's attributes raises Edge2TorchError.

Attributes:

Name	Type	Description
on	str	What was attributed onto. One of "features" or "nodes".
table	DataFrame	Feature-attribution table (examples × features). Available only when on="features". Values are signed Captum scores.
sites	dict[str, DataFrame]	Per-site node-attribution tables, keyed by layer_* or step_*. Each table has the visible graph nodes exposed at that site. Available only when on="nodes". Feedforward sites omit input nodes; state-update sites include them. Values are signed Captum scores.
summary	(nodes="hidden", aggregation="peak") -> pandas.DataFrame	Collapsed examples × nodes table from sites. nodes is "hidden", "non_input", or "all". "all" includes inputs only on state_update. aggregation is "peak", "mean_abs", or "last" and applies only to state_update. "peak" and "last" keep the sign; "mean_abs" does not. Feedforward concatenates disjoint columns and warns if a non-default aggregation is passed.