API reference

Every name below is imported straight from the package (from kpnn2 import ...). Module paths inside kpnn2 are private and may change between releases.

Callables

Names you call or construct.

parse_layered

parse_layered(edgelist: DataFrame) -> LayeredSpec

Parse a source/target edgelist into a LayeredSpec.

The graph must be a DAG. Nodes are ranked with Kahn's algorithm: input nodes (in-degree 0) have depth 0, and every other node has depth = 1 + max(parent depths). Names are sorted alphabetically inside each layer. That ranking defines LayeredSpec.layer_nodes and one Hop per layer after the first.

Every edge lands in exactly one hop mask, the one of its target layer, whether its depth gap is 1 or larger. A hop whose target has parents further back reads several layers: its mask columns are those layers concatenated. Edges with a gap greater than 1 are additionally listed in skips as metadata, so they can be reported, but they are not a separate computation and are not expanded into dummy neurons.

Terminals that are not at maximum depth (early outputs) are allowed. Cycles, self-loops, and graphs with no input or no output are not. Isolated nodes cannot appear: the node set is the union of source and target values only.

Parameters:

Name	Type	Description	Default
edgelist	DataFrame	Edge table with required columns source and target. Node names are stored as strings via str(...). Extra columns are ignored.	required

Returns:

Type	Description
LayeredSpec	Frozen structure: layers, hops, and skip metadata.

Raises:

Type	Description
Kpnn2Error	If edgelist is not a DataFrame; required columns are missing; the table is empty; names are missing or empty; source-target pairs are duplicated (message names the unique pairs, sorted); any edge is a self-loop (message names the unique nodes, sorted); the graph has a cycle (message names unranked leftover nodes); or there is no in-degree-0 node or no out-degree-0 node.

Notes

A cycle is detected when the Kahn sweep leaves some nodes unranked. The error still says the edgelist has a cycle and that only DAGs are supported, and lists every leftover name (nodes minus keys of depths), sorted alphabetically. That set may include nodes downstream of a cycle, not only vertices on a directed cycle.

Duplicate (source, target) pairs name the unique pairs as {source} -> {target}, sorted lexicographically. Self-loops name the unique nodes, sorted alphabetically.

len(hops) is len(layer_nodes) - 1 and hops[i].target_layer is i + 1. hops[i].mask has shape (layer_dims[i + 1], sum(hops[i].source_dims)), matching nn.Linear.weight, and dtype float32. Its rows are layer_nodes[i + 1] and its columns are hops[i].source_nodes, the source layers concatenated in ascending order. An entry is 1.0 only for an original edge between the node naming that row and the node naming that column. hops[0] always reads layer 0 alone, so its mask is what an align_inputs tensor feeds directly.

Every original edge with depth gap greater than 1 appears once in skips. Each record has source, target, source_layer, target_layer, source_index, and target_index. Adjacent edges never appear in skips. Membership in skips changes nothing about how the edge is computed; it is already in its target's hop mask.

Examples:

A chain plus one skip A -> C:

>>> import pandas as pd
>>> import kpnn2
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["A", "H", "A"],
...         "target": ["H", "C", "C"],
...     }
... )
>>> spec = kpnn2.parse_layered(edgelist)
>>> spec.input_nodes
('A',)
>>> spec.hidden_nodes
('H',)
>>> spec.output_nodes
('C',)
>>> spec.layer_nodes
(('A',), ('H',), ('C',))
>>> spec.layer_dims
(1, 1, 1)

The hop into C reads both earlier layers, so the skip A -> C is a column of its mask:

>>> spec.hops[1].source_layers
(0, 1)
>>> spec.hops[1].source_nodes
('A', 'H')
>>> spec.hops[1].mask.tolist()
[[1.0, 1.0]]
>>> spec.skips[0].source, spec.skips[0].target
('A', 'C')
>>> spec.skips[0].source_layer, spec.skips[0].target_layer
(0, 2)

parse_adjacency

parse_adjacency(edgelist: DataFrame) -> AdjacencySpec

Parse a source/target edgelist into an AdjacencySpec.

Every node goes into one state vector, sorted alphabetically, and every edge goes into packed source/target index tuples. Nothing is ranked, so cycles and self-loops are allowed. Use this layout for recurrent networks; use parse_layered for a DAG that should become one mask per layer.

A DAG is valid input to both parsers. The layout is a choice, not a property of the graph, so this function never inspects the graph to decide which spec to return.

Isolated nodes cannot appear: the node set is the union of source and target values only. Graphs with no in-degree-0 node or no out-degree-0 node are rejected, which also rejects a pure ring and a lone self-loop.

Parameters:

Name	Type	Description	Default
edgelist	DataFrame	Edge table with required columns source and target. Node names are stored as strings via str(...). Extra columns are ignored.	required

Returns:

Type	Description
AdjacencySpec	Frozen structure: node names, packed edge indices, and the input and output positions in the state vector.

Raises:

Type	Description
Kpnn2Error	If edgelist is not a DataFrame; required columns are missing; the table is empty; names are missing or empty; source-target pairs are duplicated (message names the unique pairs, sorted); or there is no in-degree-0 node or no out-degree-0 node.

See Also

parse_layered : Rank a DAG into one incoming mask per layer.

Notes

Self-loops are allowed here and rejected by parse_layered. That is the only edgelist rule the two parsers disagree on; every other validation is shared, so the messages match.

A self-loop removes its node from both the input set and the output set, so an edgelist of only A -> A raises for having no input node. A pure ring such as A -> B, B -> A raises for the same reason.

This function does not allocate an (n, n) tensor. Packed indices have the same length as the edge count. Order is canonical: lexicographic by (source name, target name). A dense square would have 1.0 at [target_index[i], source_index[i]], matching the nn.Linear.weight layout used by the layered hop masks. Call spec.to_mask() to materialize that square.

This function does not build an nn.Module, unroll time, choose a step count, or re-inject inputs between steps. The recurrence stays in user forward() code.

Examples:

An input feeding a two-node feedback core plus one output:

>>> import pandas as pd
>>> import kpnn2
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["x", "a", "b", "a"],
...         "target": ["a", "b", "a", "y"],
...     }
... )
>>> spec = kpnn2.parse_adjacency(edgelist)
>>> spec.nodes
('a', 'b', 'x', 'y')
>>> spec.input_nodes
('x',)
>>> spec.input_index
(2,)
>>> spec.source_index
(0, 0, 1, 2)
>>> spec.target_index
(1, 3, 0, 0)
>>> tuple(spec.to_mask().shape)
(4, 4)

The feedback edge b -> a and the forward edge a -> b are both present:

>>> mask = spec.to_mask()
>>> mask[0, 1].item(), mask[1, 0].item()
(1.0, 1.0)

MaskedLinear

MaskedLinear(mask: Tensor, bias: bool = True)

Bases: Module

Affine hop with a fixed connectivity mask.

Same job as torch.nn.Linear: call layer(x) in an nn.Module. This is not a subclass of Linear. For shapes, bias, calling the module, and training, see the PyTorch docs for torch.nn.Linear.

The mask is applied with torch.nn.utils.parametrize.register_parametrization, so layer.weight is the effective masked weight (recomputed, not an nn.Parameter). The trainable tensor lives at layer.parametrizations.weight.original. model.parameters() includes that tensor. A param-group filter that uses "weight" in name matches it; name.endswith(".weight") does not.

Parameters:

Name	Type	Description	Default
mask	Tensor	Connectivity of shape (out_features, in_features). Non-finite values are not special-cased: the tensor is stored as float32 and multiplied with the trainable weight.	required
bias	bool	If True, learn a bias of shape (out_features,). If False, there is no bias.	True

Attributes:

Name	Type	Description
in_features	int	Number of input columns, mask.shape[1].
out_features	int	Number of output columns, mask.shape[0].
weight	Tensor	Effective weight, the product of parametrizations.weight.original and a mask cast to that tensor's dtype and device. Recomputed on every access, so it is not a parameter: in-place writes to it are discarded. Assigning (layer.weight = w, under torch.no_grad()) copies w into original, where the mask hides the entries it zeroes.
parametrizations	ModuleDict	Holds parametrizations.weight.original, the trainable nn.Parameter of shape (out_features, in_features), and the mask module. Masked-out entries can be nonzero in original; they never reach the output.
mask	Tensor	Float32 buffer, same shape as the constructor mask. Not trained and not saved in state_dict. An independent copy of the constructor tensor (including spec.hops[i].mask), so later edits to that tensor do not reach this layer. Stays float32 after .half() / bfloat16 / .double(). Treat it as read-only: like any PyTorch buffer it can be written to, and doing so silently rewires the layer. Rebuild from the edgelist instead.
bias	Parameter | None	Trainable bias, or None when constructed with bias=False.

Raises:

Type	Description
Kpnn2Error	If mask is not a torch.Tensor, or is not 2-D.

Notes

Construct with mask; sizes come from mask.shape. mask is an ordinary float32 buffer: not trained, omitted from state_dict, and not a tensor subclass, so nothing custom runs per operation and torch.compile sees plain tensors. Nothing prevents writing to it; treat it as read-only and rebuild from the edgelist to change wiring. Module dtype casts do not change the stored mask dtype.

Forward is Y = F.linear(X, self.weight, self.bias), and self.weight is original * mask.to(dtype=original.dtype, device=original.device), so .half(), bfloat16, and .double() work like nn.Linear. There are no extra edge-weight constraints beyond the mask.

Registering a parametrization is what PyTorch does for "the effective weight is a function of a stored parameter", and it comes with that machinery's conventions:

• state_dict keys are parametrizations.weight.original, optional bias, and mask_digest; mask stays out of it. "weight" in name matches that parameter key; name.endswith(".weight") misses it. Default Adam(model.parameters()) needs no filter. mask_digest is a 1-D CPU uint8 tensor of length 32: the SHA-256 of the live mask's float32 C-contiguous bytes at save time, not a registered buffer. load_state_dict raises Kpnn2Error when a present digest does not match this layer's mask, and does not load the weights. A missing digest is not an error, even with strict=True. The digest catches same-shape rewiring, not a rename that leaves the 0/1 pattern unchanged (that is spec.fingerprint).
• repr reports ParametrizedMaskedLinear, because PyTorch swaps in a subclass to install the weight property. isinstance(layer, MaskedLinear) is still True.
• copy.deepcopy works; pickling the module object does not, here as for any parametrized module. Save state_dict, not the module.
• Utilities that need weight to be a raw nn.Parameter, such as torch.nn.utils.prune, reject a parametrized weight here exactly as they do on a parametrized nn.Linear. Point them at layer.parametrizations.weight under the name original; their mask then composes with the connectivity mask.
• Do not call parametrize.remove_parametrizations on weight: that drops the mask and leaves a dense layer.

reset_parameters uses per-row mask degree as fan_in, not full in_features. Typical construction: MaskedLinear(spec.hops[i].mask). Because a hop mask carries every parent of its target, including skip parents, that per-row degree is the unit's real fan-in.

Examples:

One hop whose second output is connected only to the second input:

>>> import torch
>>> import kpnn2
>>> mask = torch.tensor(
...     [
...         [1.0, 1.0],
...         [0.0, 1.0],
...     ]
... )
>>> layer = kpnn2.MaskedLinear(
...     mask,
...     bias=False,
... )
>>> layer.in_features, layer.out_features
(2, 2)
>>> tuple(layer.mask.shape)
(2, 2)
>>> x = torch.ones(3, 2)
>>> y = layer(x)
>>> tuple(y.shape)
(3, 2)

layer.weight is the masked product; the trainable tensor is one level down:

>>> bool(layer.weight[1, 0] == 0.0)
True
>>> tuple(layer.parametrizations.weight.original.shape)
(2, 2)

A zero in the mask blocks that input column:

>>> mask = torch.tensor([[1.0, 0.0]])
>>> layer = kpnn2.MaskedLinear(
...     mask,
...     bias=False,
... )
>>> a = layer(torch.tensor([[1.0, 0.0]]))
>>> b = layer(torch.tensor([[1.0, 99.0]]))
>>> torch.equal(a, b)
True

reset_parameters

reset_parameters() -> None

Initialize from per-row mask degree, not full width.

This is the difference from torch.nn.Linear.reset_parameters. For output row j, fan_in is the number of ones in mask[j]. That row of parametrizations.weight.original (and bias[j], if present) is drawn uniformly from [-1 / sqrt(fan_in), 1 / sqrt(fan_in)]. If fan_in == 0, the row and bias entry stay 0.

Degrees are counted in one pass over mask, so no per-row device synchronization happens here. Rows are then drawn one at a time, which writes straight into the trainable tensor without a full-size temporary.

forward

forward(x: Tensor) -> torch.Tensor

F.linear of x with the effective weight.

self.weight comes from the mask parametrization: the trainable tensor times a mask cast to its dtype and device. The stored mask remains float32, so .half(), bfloat16, and .double() match nn.Linear.

PackedLinear

PackedLinear(
    source_index: object,
    target_index: object,
    out_features: int,
    in_features: int,
    bias: bool = True,
)

Bases: Module

Affine map with one trainable scalar per live edge.

Packed 1-D weights, one per live edge; not torch.sparse; forward is index_add. This is not a subclass of Linear and not a full model. Dead edges are omitted: there is no dense (out_features, in_features) parameter and forward never allocates that square.

Parameters:

Name	Type	Description	Default
source_index	torch.Tensor or sequence of int	1-D integer indices of length nnz >= 1. Entry i is the input column of live edge i. Copied to an int64 buffer.	required
target_index	torch.Tensor or sequence of int	1-D integer indices of the same length. Entry i is the output row of live edge i. Copied to an int64 buffer.	required
out_features	int	Width of the output axis. Must be a positive int.	required
in_features	int	Width of the input axis. Must be a positive int.	required
bias	bool	If True, learn a bias of shape (out_features,). If False, there is no bias.	True

Attributes:

Name	Type	Description
in_features	int	Number of input columns.
out_features	int	Number of output columns.
nnz	int	Number of live edges, weight.shape[0].
weight	Parameter	Trainable packed weights of shape (nnz,). One scalar per live edge, in the same order as the index buffers. This is an ordinary parameter; there is no parametrize and no dense (out, in) tensor.
source_index	Tensor	Int64 buffer of input columns, length nnz.
target_index	Tensor	Int64 buffer of output rows, length nnz.
bias	Parameter | None	Trainable bias, or None when constructed with bias=False.

Raises:

Type	Description
Kpnn2Error	If the indices are empty, not 1-D integers, mismatched in length, out of range, duplicated as (source, target) pairs, or if out_features / in_features are not positive ints.

Notes

Construct from packed indices, not from a dense mask and not from an AdjacencySpec::

PackedLinear(
    spec.source_index,
    spec.target_index,
    len(spec.nodes),
    len(spec.nodes),
)

0 <= source_index < in_features and 0 <= target_index < out_features. Duplicate pairs are rejected. Index buffers stay integer after .half() / bfloat16 / .double(); weight and bias follow the module floating dtype like nn.Linear.

Forward gathers x[..., source_index], multiplies by weight, and index_adds into a zeros tensor of shape (..., out_features). It does not scatter into a dense (out, in) matrix and does not import torch.sparse.

reset_parameters uses per-row packed degree as fan_in, counted with bincount over target_index. Each live edge into row j is drawn uniformly from [-1 / sqrt(fan_in), 1 / sqrt(fan_in)]. If fan_in == 0, that row has no packed weights and bias[j] stays 0. Input nodes on an AdjacencySpec have in-degree 0, so they have no packed incoming edges; this layer does not invent identity connections.

state_dict keys are weight, optional bias, source_index, target_index, and index_digest. index_digest is a 1-D CPU uint8 tensor of length 32: the SHA-256 of the live index buffers' int64 C-contiguous bytes plus out_features and in_features as fixed-width integers, not a registered persistent buffer. load_state_dict raises Kpnn2Error when a present digest does not match this layer, and does not load the weights. A missing digest is not an error, even with strict=True. copy.deepcopy works.

Typical construction for a large AdjacencySpec is this class; MaskedLinear(spec.to_mask()) remains the dense path for small graphs.

Examples:

Two crossed edges on a 2-wide state, no bias:

>>> import torch
>>> import kpnn2
>>> layer = kpnn2.PackedLinear(
...     [0, 1],
...     [1, 0],
...     2,
...     2,
...     bias=False,
... )
>>> layer.in_features, layer.out_features, layer.nnz
(2, 2, 2)
>>> x = torch.ones(3, 2)
>>> y = layer(x)
>>> tuple(y.shape)
(3, 2)

reset_parameters

reset_parameters() -> None

Initialize from per-row packed degree, not full width.

For output row j, fan_in is the number of packed edges with target_index == j (bincount, minlength=out_features). Each live edge into that row, and bias[j] if present, is drawn uniformly from [-1 / sqrt(fan_in), 1 / sqrt(fan_in)]. If fan_in == 0, that row has no packed weights and bias[j] stays 0.

forward

forward(x: Tensor) -> torch.Tensor

Gather live inputs, scale by packed weights, index_add.

contrib = x[..., source_index] * weight, then index_add into zeros of shape (..., out_features). Adds bias when present. Packed 1-D weights, one per live edge; not torch.sparse; forward is index_add.

gather_hop_inputs

gather_hop_inputs(
    saved: Mapping[int, Tensor], hop: Hop
) -> torch.Tensor

Concatenate the layer tensors one hop reads, in mask-column order.

Call this in forward() just before MaskedLinear(hop.mask). It sits between hops. It does not inject values into the previous layer, does not pick skip nodes by name, and holds no weights.

A hop mask's columns are whole source layers laid side by side (hop.source_layers). This function builds that tensor from saved:

• Adjacent hop (no skips): one source, the layer below. That saved tensor is returned as-is, with no copy.
• Hop with skips: the previous layer plus older layers, concatenated on the last axis. The hop mask, not this gather, zeros columns that are not edges. Example: skip A → C with [A, B] then [H] yields [A, B, H].

Store every layer you produce in saved. A missing source layer (not a missing node) raises Kpnn2Error instead of silently dropping those edges. Unused keys are ignored; saved tensors are not modified.

Parameters:

Name	Type	Description	Default
saved	mapping of int to torch.Tensor	Layer index to that layer's activation. Width of saved[i] must be layer_dims[i]. Only the layers in hop.source_layers are read.	required
hop	Hop	The hop about to be applied, from spec.hops.	required

Returns:

Type	Description
Tensor	Shape (..., hop.mask.shape[1]), ready for MaskedLinear(hop.mask).

Raises:

Type	Description
Kpnn2Error	If saved is not a mapping or hop is not a Hop; a needed layer is missing from saved or is not a tensor; a saved tensor has the wrong number of units; or the parts disagree on dtype or device.

Examples:

The hop into C reads layers 0 and 1, so its input is two columns wide:

>>> import pandas as pd
>>> import torch
>>> import kpnn2
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["A", "H", "A"],
...         "target": ["H", "C", "C"],
...     }
... )
>>> spec = kpnn2.parse_layered(edgelist)
>>> saved = {
...     0: torch.tensor([[2.0]]),
...     1: torch.tensor([[5.0]]),
... }
>>> x = kpnn2.gather_hop_inputs(
...     saved,
...     spec.hops[1],
... )
>>> x.tolist()
[[2.0, 5.0]]

Forgetting to store a layer raises instead of silently dropping the edges that read it:

>>> kpnn2.gather_hop_inputs(
...     {1: torch.tensor([[5.0]])},
...     spec.hops[1],
... )
Traceback (most recent call last):
...
Kpnn2Error: saved is missing layer 0. ...

align_inputs

align_inputs(
    data: DataFrame, spec: LayeredSpec | AdjacencySpec
) -> torch.Tensor

Return a float32 tensor whose columns follow spec.input_nodes.

DataFrame. Required columns are spec.input_nodes. Labels are matched after str(...), the same conversion used for edgelist node names, so an integer column 1 matches node "1". Extra columns are ignored. Columns are reordered to spec.input_nodes. Missing, duplicate, or non-numeric required columns raise an error.

Tensor. Not accepted. Raise Kpnn2Error. Pre-ordered tensors go straight to the model. Users who need alignment pass a DataFrame.

AnnData is not supported.

Parameters:

Name	Type	Description	Default
data	DataFrame	Feature table with named columns.	required
spec	LayeredSpec or AdjacencySpec	Graph structure whose input_nodes define column order. Both spec types work the same way here; only input_nodes is read.	required

Returns:

Type	Description
Tensor	Float32 tensor of shape (n_samples, len(spec.input_nodes)).

Raises:

Type	Description
Kpnn2Error	If spec is neither a LayeredSpec nor an AdjacencySpec; data is a tensor; data is not a DataFrame; required DataFrame columns are missing or duplicated (including after str conversion; the message names the unique duplicated labels, sorted, comma-separated); or required columns are non-numeric.

Notes

PyTorch never sees feature names. Column i of the returned tensor is spec.input_nodes[i]. Use this function when the table has named columns. A tensor whose columns already follow spec.input_nodes goes straight to the model. Passing DataFrame.to_numpy() (or any hand-stacked array) into the model can silently wire the wrong features if the column order differs.

The returned width is always len(spec.input_nodes). For a LayeredSpec that is the width of hops[0].mask, whose only source layer is layer 0, so the tensor feeds the first hop directly and needs no gathering. For an AdjacencySpec it is not the state width: scatter the tensor into the len(spec.nodes)-wide state vector with spec.input_index before calling MaskedLinear(spec.to_mask()).

Examples:

Extra columns are dropped and remaining columns are reordered:

>>> import pandas as pd
>>> import torch
>>> import kpnn2
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["A", "H"],
...         "target": ["H", "C"],
...     }
... )
>>> spec = kpnn2.parse_layered(edgelist)
>>> spec.input_nodes
('A',)
>>> df = pd.DataFrame(
...     {
...         "unused": [9.0, 8.0],
...         "A": [0.5, 1.5],
...     }
... )
>>> x = kpnn2.align_inputs(df, spec)
>>> x.dtype
torch.float32
>>> tuple(x.shape)
(2, 1)
>>> x.tolist()
[[0.5], [1.5]]

An AdjacencySpec works the same way, but the result is len(input_nodes) wide and must be scattered into the state vector before it reaches MaskedLinear(spec.to_mask()):

>>> cyclic = pd.DataFrame(
...     {
...         "source": ["x", "a", "b", "a"],
...         "target": ["a", "b", "a", "y"],
...     }
... )
>>> state_spec = kpnn2.parse_adjacency(cyclic)
>>> inputs = pd.DataFrame({"x": [0.5, 1.5]})
>>> x = kpnn2.align_inputs(inputs, state_spec)
>>> tuple(x.shape), tuple(state_spec.to_mask().shape)
((2, 1), (4, 4))
>>> state = torch.zeros(
...     2,
...     len(state_spec.nodes),
... )
>>> state[:, state_spec.input_index] = x
>>> state.tolist()
[[0.0, 0.0, 0.5, 0.0], [0.0, 0.0, 1.5, 0.0]]

A tensor is not accepted; pass a DataFrame instead:

>>> t = torch.tensor([[0.5], [1.5]])
>>> kpnn2.align_inputs(t, spec)
Traceback (most recent call last):
...
Kpnn2Error: 'data' is a tensor; a pandas DataFrame is required. ...

map_node_attributions

map_node_attributions(
    attributions: Tensor | Sequence[Tensor],
    spec: LayeredSpec | AdjacencySpec,
    layer: int | None = None,
    *,
    dims: Sequence[str] | None = None,
    coords: Mapping[str, Sequence] | None = None
) -> xr.DataArray

Label an attribution tensor with node names from a spec.

Values are copied with detach() onto CPU. This function does not run an attribution method and does not import Captum. Pass the tensor (or per-call tensors) you already computed. Nothing is summed or averaged.

A 2-D (batch, n_units) tensor becomes dims (observation, node). Extra Captum axes need dims so one axis is named node. A tuple of equal-shaped tensors is stacked on a new step axis (one entry per module call).

Where the names come from depends on the spec:

• LayeredSpec: layer is required and names come from spec.layer_nodes[layer]. The result carries a scalar layer coordinate.
• AdjacencySpec: layer must be omitted and names come from spec.nodes, the whole state vector. The result carries no layer coordinate, because there is no depth.

Parameters:

Name	Type	Description	Default
attributions	Tensor or sequence of Tensor	Scores whose node axis length equals the number of named units: len(spec.layer_nodes[layer]) for a LayeredSpec, len(spec.nodes) for an AdjacencySpec. A sequence is stacked along step.	required
spec	LayeredSpec or AdjacencySpec	Graph structure supplying the node coordinate.	required
layer	int	0-based index into spec.layer_nodes. Layer 0 is the input layer. Stored as a scalar coordinate layer. Required for a LayeredSpec; must be omitted for an AdjacencySpec.	None
dims	sequence of str	Name of each axis of the (stacked) tensor. Must contain node exactly once. Required when the tensor has 3 or more axes (except a default-stacked 2-D sequence).	None
coords	mapping	Labels for axes other than node and layer. Length of each entry must match that axis. node is set from spec, and layer from the layer argument.	None

Returns:

Type	Description
DataArray	Raw scores with a node coordinate from the spec, plus a scalar layer coordinate for a LayeredSpec. Use .to_dataframe(name="score").reset_index() for a long table, or .to_pandas() for a 2-D wide table.

Raises:

Type	Description
Kpnn2Error	If spec is neither a LayeredSpec nor an AdjacencySpec; layer is missing for a LayeredSpec, given for an AdjacencySpec, or not a valid int; attributions is not a tensor or a non-empty sequence of equal-shaped tensors; dims is missing or inconsistent; a node axis length does not match the named units; or coords is invalid.

Notes

For a MaskedLinear built from spec.hops[i].mask, the layer index to pass here is spec.hops[i].target_layer, that is i + 1: the hop output, not its input. Do not name-map tensors from BatchNorm or other unnamed modules; only map units that are spec nodes.

A recurrent net built on an AdjacencySpec has no layers to index. The natural extra axis there is step: pass one tensor per time step as a sequence and they are stacked for you.

Examples:

Name a two-row tensor at the output layer:

>>> import pandas as pd
>>> import torch
>>> import kpnn2
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["A", "H"],
...         "target": ["H", "C"],
...     }
... )
>>> spec = kpnn2.parse_layered(edgelist)
>>> spec.layer_nodes
(('A',), ('H',), ('C',))
>>> scores = torch.tensor([[0.5], [1.0]])
>>> da = kpnn2.map_node_attributions(
...     attributions=scores,
...     spec=spec,
...     layer=2,
... )
>>> da["node"].values.tolist()
['C']
>>> da.sel(node="C").values.tolist()
[0.5, 1.0]
>>> int(da.coords["layer"])
2

The layer argument is an index into spec.layer_nodes:

>>> hidden = kpnn2.map_node_attributions(
...     attributions=torch.zeros(2, 1),
...     spec=spec,
...     layer=1,
... )
>>> hidden["node"].values.tolist()
['H']

On an AdjacencySpec there are no layers: omit layer and the whole state vector is named. One tensor per time step stacks onto a step axis:

>>> cyclic = pd.DataFrame(
...     {
...         "source": ["x", "a", "b", "a"],
...         "target": ["a", "b", "a", "y"],
...     }
... )
>>> state_spec = kpnn2.parse_adjacency(cyclic)
>>> per_step = kpnn2.map_node_attributions(
...     attributions=[
...         torch.zeros(2, 4),
...         torch.ones(2, 4),
...     ],
...     spec=state_spec,
... )
>>> per_step.dims
('step', 'observation', 'node')
>>> per_step["node"].values.tolist()
['a', 'b', 'x', 'y']
>>> "layer" in per_step.coords
False

Specs

Returned by the parsers. Frozen dataclasses; treat masks as read-only.

LayeredSpec dataclass

LayeredSpec(
    input_nodes: tuple[str, ...],
    output_nodes: tuple[str, ...],
    hidden_nodes: tuple[str, ...],
    layer_nodes: tuple[tuple[str, ...], ...],
    layer_dims: tuple[int, ...],
    hops: tuple[Hop, ...],
    skips: tuple[Skip, ...],
)

Frozen blueprint from parse_layered.

Structure only: not an nn.Module and no parameters. One Hop per layer after the first, and one MaskedLinear per hop is the whole model wiring.

Parameters:

Name	Type	Description	Default
input_nodes	tuple[str, ...]	In-degree 0 names, alphabetical. This is the column order of tensors returned by align_inputs.	required
output_nodes	tuple[str, ...]	Out-degree 0 names, alphabetical.	required
hidden_nodes	tuple[str, ...]	Names that are neither input nor output, alphabetical.	required
layer_nodes	tuple[tuple[str, ...], ...]	layer_nodes[i] is the names at depth i, alphabetical. Index 0 is the input layer.	required
layer_dims	tuple[int, ...]	layer_dims[i] == len(layer_nodes[i]).	required
hops	tuple[Hop, ...]	One hop per layer after the first: len(hops) == len(layer_nodes) - 1 and hops[i].target_layer == i + 1. hops[i].mask holds every edge entering layer i + 1, adjacent and skip alike, over the concatenated source layers. hops[0] always reads layer 0 only.	required
skips	tuple[Skip, ...]	Original edges with depth gap greater than 1, as metadata. Each one is already a one in hops[target_layer - 1].mask; this list only says which edges span layers.	required

Notes

Fields cannot be reassigned and sequences are tuples, so the structure itself is fixed. The mask tensors are plain float32 tensors and are not write-protected; treat them as read-only and rebuild from the edgelist to change wiring. MaskedLinear clones the mask into a non-persistent buffer independent of spec.hops[i].mask, so a layer built earlier keeps its own connectivity either way.

to_edgelist() returns the original edges as a two-column source / target DataFrame, rows sorted lexicographically. parse_layered on that table reconstructs the same node lists, hops, and hop masks.

to_dict() returns a JSON-safe tagged dict (kpnn2_spec, layout, edges). from_dict rebuilds this spec by calling parse_layered. fingerprint is the SHA-256 of that canonical JSON. Pickle / torch.save of the dataclass is not the supported interchange.

Because a hop mask carries every parent of its target, the per-row degree MaskedLinear initializes from is the real fan-in of that unit, skips included.

Examples:

Inspect layers, a hop, and a skip after parsing:

>>> import pandas as pd
>>> import kpnn2
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["A", "H", "A"],
...         "target": ["H", "C", "C"],
...     }
... )
>>> spec = kpnn2.parse_layered(edgelist)
>>> spec.layer_nodes
(('A',), ('H',), ('C',))
>>> spec.layer_dims
(1, 1, 1)
>>> spec.hops[0].source_layers, spec.hops[0].mask.tolist()
((0,), [[1.0]])
>>> spec.hops[1].source_layers, spec.hops[1].mask.tolist()
((0, 1), [[1.0, 1.0]])
>>> spec.skips[0].source, spec.skips[0].target
('A', 'C')
>>> spec.skips[0].source_layer, spec.skips[0].target_layer
(0, 2)

fingerprint property

fingerprint: str

SHA-256 hex digest of the canonical to_dict() JSON.

The payload is json.dumps(self.to_dict(), sort_keys=True, separators=(",", ":"), ensure_ascii=False) encoded as UTF-8. The result is 64 lowercase hex characters. It is not Python hash().

Returns:

Type	Description
str	Hex digest of the tagged spec dict.

Examples:

>>> import pandas as pd
>>> import kpnn2
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["A", "H"],
...         "target": ["H", "C"],
...     }
... )
>>> spec = kpnn2.parse_layered(edgelist)
>>> len(spec.fingerprint)
64
>>> (
...     spec.fingerprint
...     == kpnn2.parse_layered(spec.to_edgelist()).fingerprint
... )
True

to_edgelist

to_edgelist() -> pd.DataFrame

Return this spec's edges as a two-column table.

Columns are exactly source then target. Rows follow the hop masks in canonical order: sorted lexicographically by (source, target), one row per original edge, names as strings. Extra columns from the DataFrame that was parsed are not reproduced.

parse_layered on this table reconstructs the same node lists, hops, and hop masks. Skip tuple order follows these sorted rows rather than the original parse input order; the skip set matches.

Returns:

Type	Description
DataFrame	One row per original edge.

Examples:

Unsorted input comes back sorted:

>>> import pandas as pd
>>> import kpnn2
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["H", "A", "A"],
...         "target": ["C", "C", "H"],
...     }
... )
>>> spec = kpnn2.parse_layered(edgelist)
>>> table = spec.to_edgelist()
>>> list(table.columns)
['source', 'target']
>>> table["source"].tolist()
['A', 'A', 'H']
>>> table["target"].tolist()
['C', 'H', 'C']

to_dict

to_dict() -> dict

Return this spec as a JSON-safe tagged dict.

Keys are kpnn2_spec (integer 1), layout ("layered"), and edges (list of [source, target] lists in the same order as to_edgelist() rows). The returned dict is new on every call.

Returns:

Type	Description
dict	Tagged edge list plus layout.

Examples:

>>> import pandas as pd
>>> import kpnn2
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["A", "H"],
...         "target": ["H", "C"],
...     }
... )
>>> spec = kpnn2.parse_layered(edgelist)
>>> payload = spec.to_dict()
>>> payload["kpnn2_spec"]
1
>>> payload["layout"]
'layered'
>>> payload["edges"]
[['A', 'H'], ['H', 'C']]

from_dict classmethod

from_dict(payload: dict) -> LayeredSpec

Rebuild a LayeredSpec from to_dict() output.

Calls parse_layered on a DataFrame built from payload["edges"]. Hops and masks are not assembled by hand. Extra unknown keys are ignored.

Parameters:

Name	Type	Description	Default
payload	dict	A dict with kpnn2_spec, layout, and edges. layout must be "layered".	required

Returns:

Type	Description
LayeredSpec	The parsed spec.

Raises:

Type	Description
Kpnn2Error	If payload is not a dict; kpnn2_spec is missing or not 1; layout is missing, not a known layout, or is "adjacency"; or edges is missing or not a sequence of two nonempty names.

Examples:

>>> import pandas as pd
>>> import kpnn2
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["A", "H"],
...         "target": ["H", "C"],
...     }
... )
>>> spec = kpnn2.parse_layered(edgelist)
>>> roundtrip = kpnn2.LayeredSpec.from_dict(spec.to_dict())
>>> roundtrip.layer_nodes == spec.layer_nodes
True

Hop dataclass

Hop(
    target_layer: int,
    source_layers: tuple[int, ...],
    source_dims: tuple[int, ...],
    source_nodes: tuple[str, ...],
    mask: Tensor,
)

Every edge entering one layer, as one mask.

A hop is what a single MaskedLinear computes. Its mask covers all parents of target_layer, whether they sit in the layer directly below or several layers back, so a skip edge is an ordinary one in this mask rather than a separate term added later. That is what makes an edge impossible to lose: apply the hop and every parent is applied with it.

The mask columns are the source layers concatenated in ascending order, which is the axis kpnn2.gather_hop_inputs builds.

Parameters:

Name	Type	Description	Default
target_layer	int	Depth of the layer this hop produces. Always at least 1; layer 0 has no parents.	required
source_layers	tuple[int, ...]	Depths this hop reads, ascending, each one below target_layer. Only layers that really feed the target appear, and target_layer - 1 is always one of them. A hop with a single entry is a plain adjacent hop.	required
source_dims	tuple[int, ...]	Units contributed by each entry of source_layers, same order. Their sum is mask.shape[1].	required
source_nodes	tuple[str, ...]	Node names of the mask columns, source layers concatenated in source_layers order. One name per node; with one unit per node that is one name per column.	required
mask	Tensor	Connectivity of shape (layer_dims[target_layer], sum(source_dims)), dtype float32, matching nn.Linear.weight. Entry [target_index, column] is 1.0 for an original edgelist edge and 0.0 otherwise. Treat it as read-only: it is a plain tensor, so writing to it silently changes the wiring this record describes.	required

Notes

To locate one source layer's block inside the mask, add the widths in front of it:

offset = sum(source_dims[:source_layers.index(layer)])

column_offsets does that for you.

Examples:

A chain A -> H -> C plus the skip A -> C:

>>> import pandas as pd
>>> import kpnn2
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["A", "H", "A"],
...         "target": ["H", "C", "C"],
...     }
... )
>>> spec = kpnn2.parse_layered(edgelist)
>>> hop = spec.hops[1]
>>> hop.target_layer, hop.source_layers
(2, (0, 1))
>>> hop.source_nodes
('A', 'H')
>>> hop.mask.tolist()
[[1.0, 1.0]]

column_offsets property

column_offsets: tuple[int, ...]

First mask column of each entry of source_layers.

Same length and order as source_layers. Add a node's index inside its own layer to get its mask column.

Skip dataclass

Skip(
    source: str,
    target: str,
    source_layer: int,
    target_layer: int,
    source_index: int,
    target_index: int,
)

One original edge whose endpoints are more than one layer apart.

This is metadata, not a separate computation. The edge itself is a one in LayeredSpec.hops[target_layer - 1].mask, exactly like an adjacent edge, so nothing has to add it back later and nothing can forget to. Read skips to report or inspect which prior-knowledge edges span layers; do not expand them into dummy neurons.

Parameters:

Name	Type	Description	Default
source	str	Source node name.	required
target	str	Target node name.	required
source_layer	int	Depth of source in LayeredSpec.layer_nodes.	required
target_layer	int	Depth of target. Always satisfies target_layer - source_layer > 1.	required
source_index	int	Column index of source in layer_nodes[source_layer].	required
target_index	int	Column index of target in layer_nodes[target_layer].	required

AdjacencySpec dataclass

AdjacencySpec(
    nodes: tuple[str, ...],
    input_nodes: tuple[str, ...],
    output_nodes: tuple[str, ...],
    hidden_nodes: tuple[str, ...],
    source_index: tuple[int, ...],
    target_index: tuple[int, ...],
    input_index: tuple[int, ...],
    output_index: tuple[int, ...],
)

Frozen blueprint from parse_adjacency.

Structure only: not an nn.Module and no parameters. Every node lives in one state vector and connectivity is packed as source/target index tuples, so the graph may contain cycles and self-loops. There is no stored square mask.

There are no depths here: no layer_nodes, no hops tuple, and no skips. This is not a one-layer LayeredSpec. Use PackedLinear on source_index / target_index for a state update that never allocates (n, n), or MaskedLinear(spec.to_mask()) for the dense path, and write the recurrence in forward().

Parameters:

Name	Type	Description	Default
nodes	tuple[str, ...]	Every node name, alphabetical. This is the unit order of the state vector and the row and column order of to_mask().	required
input_nodes	tuple[str, ...]	In-degree 0 names, alphabetical. This is the column order of tensors returned by align_inputs.	required
output_nodes	tuple[str, ...]	Out-degree 0 names, alphabetical.	required
hidden_nodes	tuple[str, ...]	Names that are neither input nor output, alphabetical.	required
source_index	tuple[int, ...]	For each original edge, the column in nodes (the source). Same length as target_index and as the edge count. Order is canonical: lexicographic by (source name, target name), identical to to_edgelist() row order. Cycles and self-loops are included.	required
target_index	tuple[int, ...]	For each original edge, the row in nodes (the target). A dense square would have 1.0 at [target_index[i], source_index[i]].	required
input_index	tuple[int, ...]	Position of each input_nodes name in nodes.	required
output_index	tuple[int, ...]	Position of each output_nodes name in nodes.	required

Notes

Fields cannot be reassigned and sequences are tuples, so the structure itself is fixed. There is no mask field and no densifying mask property. to_mask() allocates a fresh dense square on every call; mutating that tensor does not change this spec. MaskedLinear(spec.to_mask()) clones the square into a non-persistent buffer, so a layer built earlier keeps its own connectivity.

align_inputs returns len(input_nodes) columns, which is not the state width. Scatter that tensor into the n-wide state vector with input_index. Input rows of to_mask() are all zeros, so under the degree-aware init of MaskedLinear they stay zero: writing the inputs in is required, not cosmetic.

to_edgelist() returns the original edges as a two-column source / target DataFrame, rows sorted lexicographically, including cycle edges and self-loops. parse_adjacency on that table reconstructs this spec's node lists, packed indices, and input/output indices.

to_dict() returns a JSON-safe tagged dict (kpnn2_spec, layout, edges). from_dict rebuilds this spec by calling parse_adjacency. fingerprint is the SHA-256 of that canonical JSON. Pickle / torch.save of the dataclass is not the supported interchange.

Examples:

An input feeding a two-node feedback core plus one output:

>>> import pandas as pd
>>> import kpnn2
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["x", "a", "b", "a"],
...         "target": ["a", "b", "a", "y"],
...     }
... )
>>> spec = kpnn2.parse_adjacency(edgelist)
>>> spec.nodes
('a', 'b', 'x', 'y')
>>> spec.input_nodes, spec.output_nodes
(('x',), ('y',))
>>> spec.hidden_nodes
('a', 'b')
>>> spec.input_index, spec.output_index
((2,), (3,))
>>> spec.source_index
(0, 0, 1, 2)
>>> spec.target_index
(1, 3, 0, 0)
>>> tuple(spec.to_mask().shape)
(4, 4)
>>> spec.to_mask()[0].tolist()
[0.0, 1.0, 1.0, 0.0]

fingerprint property

fingerprint: str

SHA-256 hex digest of the canonical to_dict() JSON.

The payload is json.dumps(self.to_dict(), sort_keys=True, separators=(",", ":"), ensure_ascii=False) encoded as UTF-8. The result is 64 lowercase hex characters. It is not Python hash().

Returns:

Type	Description
str	Hex digest of the tagged spec dict.

Examples:

>>> import pandas as pd
>>> import kpnn2
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["x", "a", "b", "a"],
...         "target": ["a", "b", "a", "y"],
...     }
... )
>>> spec = kpnn2.parse_adjacency(edgelist)
>>> len(spec.fingerprint)
64
>>> (
...     spec.fingerprint
...     == kpnn2.parse_adjacency(spec.to_edgelist()).fingerprint
... )
True

to_edgelist

to_edgelist() -> pd.DataFrame

Return this spec's edges as a two-column table.

Columns are exactly source then target. Rows follow the packed indices in canonical order: sorted lexicographically by (source, target), one row per original edge, names as strings, including cycle edges and self-loops. Extra columns from the DataFrame that was parsed are not reproduced.

parse_adjacency on this table reconstructs the same node lists, packed indices, and input/output indices.

Returns:

Type	Description
DataFrame	One row per original edge.

Examples:

A cycle comes back as sorted pairs:

>>> import pandas as pd
>>> import kpnn2
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["x", "a", "b", "a"],
...         "target": ["a", "b", "a", "y"],
...     }
... )
>>> spec = kpnn2.parse_adjacency(edgelist)
>>> table = spec.to_edgelist()
>>> list(table.columns)
['source', 'target']
>>> table["source"].tolist()
['a', 'a', 'b', 'x']
>>> table["target"].tolist()
['b', 'y', 'a', 'a']

to_mask

to_mask() -> Tensor

Allocate a dense float32 square from the packed edges.

Shape is (n, n) with n from the node layout (len(nodes) at width 1). The result starts at zeros; each live edge sets 1.0 at [target_index[i], source_index[i]]. Every call returns a fresh tensor. Mutating it does not change this spec or the next to_mask() call.

Returns:

Type	Description
Tensor	New dense connectivity square. This allocates.

Examples:

>>> import pandas as pd
>>> import kpnn2
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["x", "a", "b", "a"],
...         "target": ["a", "b", "a", "y"],
...     }
... )
>>> spec = kpnn2.parse_adjacency(edgelist)
>>> mask = spec.to_mask()
>>> tuple(mask.shape)
(4, 4)
>>> mask[0, 1].item(), mask[1, 0].item()
(1.0, 1.0)

to_dict

to_dict() -> dict

Return this spec as a JSON-safe tagged dict.

Keys are kpnn2_spec (integer 1), layout ("adjacency"), and edges (list of [source, target] lists in the same order as to_edgelist() rows, including cycle edges and self-loops). The returned dict is new on every call.

Returns:

Type	Description
dict	Tagged edge list plus layout.

Examples:

>>> import pandas as pd
>>> import kpnn2
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["x", "a", "b", "a"],
...         "target": ["a", "b", "a", "y"],
...     }
... )
>>> spec = kpnn2.parse_adjacency(edgelist)
>>> payload = spec.to_dict()
>>> payload["kpnn2_spec"]
1
>>> payload["layout"]
'adjacency'
>>> payload["edges"]
[['a', 'b'], ['a', 'y'], ['b', 'a'], ['x', 'a']]

from_dict classmethod

from_dict(payload: dict) -> AdjacencySpec

Rebuild an AdjacencySpec from to_dict() output.

Calls parse_adjacency on a DataFrame built from payload["edges"]. Packed indices are not assembled by hand. Extra unknown keys are ignored.

Parameters:

Name	Type	Description	Default
payload	dict	A dict with kpnn2_spec, layout, and edges. layout must be "adjacency".	required

Returns:

Type	Description
AdjacencySpec	The parsed spec.

Raises:

Type	Description
Kpnn2Error	If payload is not a dict; kpnn2_spec is missing or not 1; layout is missing, not a known layout, or is "layered"; or edges is missing or not a sequence of two nonempty names.

Examples:

>>> import pandas as pd
>>> import kpnn2
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["x", "a", "b", "a"],
...         "target": ["a", "b", "a", "y"],
...     }
... )
>>> spec = kpnn2.parse_adjacency(edgelist)
>>> roundtrip = kpnn2.AdjacencySpec.from_dict(spec.to_dict())
>>> roundtrip.nodes == spec.nodes
True

Errors and version

Kpnn2Error

Bases: Exception

User-facing failure from the public kpnn2 API.

Raised for invalid edgelists, illegal LayeredSpec operations, bad MaskedLinear masks, saved activations that do not match the hop they are gathered for, and input or attribution tensors that do not match the spec.

Examples:

>>> import pandas as pd
>>> import kpnn2
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["A"],
...         "target": ["A"],
...     }
... )
>>> kpnn2.parse_layered(edgelist)
Traceback (most recent call last):
...
Kpnn2Error: Edgelist contains 1 self-loop(s): A. ...

__version__ module-attribute

__version__ = '0.1.0'