Mapping attributions¶

An attribution method returns a tensor whose node axis is bare positions. map_node_attributions() attaches the spec's names — the spec being the frozen structure a parser returns — and returns an xarray.DataArray. Values are copied with detach() onto CPU. It does not run an attribution method, import Captum, or sum or average scores.

This page is the naming rule on one graph, then the tensor layouts. Fake tensors keep the mapping visible without training. For a trained model plus Captum, see Feedforward example Step 6. For an AdjacencySpec, see the Cyclic graph example (feedback loop, shared MaskedLinear), the Time-series example (sequence x_t), and the Transformer example. The signature is on the API page.

The naming rule¶

You produce the scores; this function names the node axis. Any method will do (Captum LayerConductance, input gradients, or your own), provided the tensor is as wide as the axis you name.

Captum scores modules; a LayeredSpec names layers. The module you build from a hop — everything entering one layer — reads one axis and writes another: a hop may read several layers, but it writes exactly one. Say which axis your tensor is on with exactly one argument:

• layer= indexes spec.layer_nodes. Layer 0 is the input nodes.
• hop_output=spec.hops[i] is what the module on spec.hops[i] returns: layer hops[i].target_layer, that is i + 1. Captum layer methods such as LayerConductance score this side by default, so you pass the same i you used for the module.
• hop_input=spec.hops[i] is what that module reads: the source layers gather_hop_inputs() concatenates (Captum with attribute_to_layer_input=True).

The hop module returns the pre-activation tensor, so hop_output names that linear map. The ReLU after it is a separate nn.ReLU module in self.acts, one per hop, and its output is the post-activation node state. To score that state, hook the nn.ReLU module with a method such as LayerActivation or LayerGradientXActivation. Keeping the ReLU as a module also lets DeepLift, DeepLiftShap, and layer-wise relevance propagation (LRP) find and rescale it.

Do not name-map BatchNorm or other unnamed modules. Only map tensors whose units are spec nodes.

With an AdjacencySpec, omit layer, hop_input, and hop_output. That layout has no depths, so the node axis is the whole state vector — a single vector over all nodes, spec.nodes — and the result carries no layer coordinate.

Input attributions are shorter than the state vector. Pass axis="inputs" to name spec.input_nodes instead; a wide input's name repeats once per unit. On a LayeredSpec, layer=0 gives the same names. The width is never used to guess the axis.

The graph below shows why the axis is always stated. Inputs A and B feed hidden nodes H1 and H2, which feed output C. Layers 0 and 1 both have width 2; layer 2 has width 1. Two layers sharing a width is ordinary, and it is why the function does not guess the axis from tensor.shape[-1].

In [ ]:

import pandas as pd
import torch

import kpnn2

edgelist = pd.DataFrame(
    {
        "source": ["A", "B", "A", "B", "H1", "H2"],
        "target": ["H1", "H1", "H2", "H2", "C", "C"],
    }
)
spec = kpnn2.parse_layered(edgelist)

print(
    "layer_nodes:",
    spec.layer_nodes,
)
print(
    "layer_dims:",
    spec.layer_dims,
)
print(
    "n hops:",
    len(spec.hops),
)

import pandas as pd import torch import kpnn2 edgelist = pd.DataFrame( { "source": ["A", "B", "A", "B", "H1", "H2"], "target": ["H1", "H1", "H2", "H2", "C", "C"], } ) spec = kpnn2.parse_layered(edgelist) print( "layer_nodes:", spec.layer_nodes, ) print( "layer_dims:", spec.layer_dims, ) print( "n hops:", len(spec.hops), )

layer_nodes: (('A', 'B'), ('H1', 'H2'), ('C',))
layer_dims: (2, 2, 1)
n hops: 2

The next cell builds one PackedLinear per hop and prints the two widths side by side. spec.hops[i] is the incoming hop that writes layer i + 1, so its output has the same width as spec.layer_nodes[i + 1]. This graph has no skip edges, so each hop reads only the layer below it.

In [ ]:

hop0 = kpnn2.PackedLinear(
    spec.hops[0].source_index,
    spec.hops[0].target_index,
    spec.hops[0].out_features,
    spec.hops[0].in_features,
)
hop1 = kpnn2.PackedLinear(
    spec.hops[1].source_index,
    spec.hops[1].target_index,
    spec.hops[1].out_features,
    spec.hops[1].in_features,
)
print(
    "hops[0] out_features:",
    hop0.out_features,
)
print(
    "layer 1 width:",
    spec.layer_dims[1],
)
print(
    "hops[1] out_features:",
    hop1.out_features,
)
print(
    "layer 2 width:",
    spec.layer_dims[2],
)

hop0 = kpnn2.PackedLinear( spec.hops[0].source_index, spec.hops[0].target_index, spec.hops[0].out_features, spec.hops[0].in_features, ) hop1 = kpnn2.PackedLinear( spec.hops[1].source_index, spec.hops[1].target_index, spec.hops[1].out_features, spec.hops[1].in_features, ) print( "hops[0] out_features:", hop0.out_features, ) print( "layer 1 width:", spec.layer_dims[1], ) print( "hops[1] out_features:", hop1.out_features, ) print( "layer 2 width:", spec.layer_dims[2], )

hops[0] out_features: 2
layer 1 width: 2
hops[1] out_features: 1
layer 2 width: 1

A (2, 2) tensor is a legal score matrix at layer 0 and at layer 1. The names come only from layer=.

In [ ]:

scores_2x2 = torch.tensor(
    [
        [0.10, 0.20],
        [0.30, 0.40],
    ]
)

named_inputs = kpnn2.map_node_attributions(
    attributions=scores_2x2,
    spec=spec,
    layer=0,
)
named_hidden = kpnn2.map_node_attributions(
    attributions=scores_2x2,
    spec=spec,
    layer=1,
)

print(
    "layer 0 nodes:",
    named_inputs["node"].values.tolist(),
)
print(
    "layer 1 nodes:",
    named_hidden["node"].values.tolist(),
)
print(
    "scalar layer coord (inputs):",
    int(named_inputs.coords["layer"]),
)
print(
    "scalar layer coord (hidden):",
    int(named_hidden.coords["layer"]),
)

scores_2x2 = torch.tensor( [ [0.10, 0.20], [0.30, 0.40], ] ) named_inputs = kpnn2.map_node_attributions( attributions=scores_2x2, spec=spec, layer=0, ) named_hidden = kpnn2.map_node_attributions( attributions=scores_2x2, spec=spec, layer=1, ) print( "layer 0 nodes:", named_inputs["node"].values.tolist(), ) print( "layer 1 nodes:", named_hidden["node"].values.tolist(), ) print( "scalar layer coord (inputs):", int(named_inputs.coords["layer"]), ) print( "scalar layer coord (hidden):", int(named_hidden.coords["layer"]), )

layer 0 nodes: ['A', 'B']
layer 1 nodes: ['H1', 'H2']
scalar layer coord (inputs): 0
scalar layer coord (hidden): 1

The same holds for the two sides of one hop. hops[0] reads A, B and writes H1, H2, so a (2, 2) Captum result on that module fits both sides. hop_output= gives the names the module writes, identical to layer=1. hop_input= gives the names it reads, with no layer coordinate, because a hop's input can span several depths.

In [ ]:

hop = spec.hops[0]
written = kpnn2.map_node_attributions(
    attributions=scores_2x2,
    spec=spec,
    hop_output=hop,
)
read = kpnn2.map_node_attributions(
    attributions=scores_2x2,
    spec=spec,
    hop_input=hop,
)

print(
    "hop_output nodes:",
    written["node"].values.tolist(),
)
print(
    "hop_input nodes:",
    read["node"].values.tolist(),
)
print(
    "hop_output equals layer=1:",
    written.identical(named_hidden),
)
print(
    "layer coord on hop_input:",
    "layer" in read.coords,
)

hop = spec.hops[0] written = kpnn2.map_node_attributions( attributions=scores_2x2, spec=spec, hop_output=hop, ) read = kpnn2.map_node_attributions( attributions=scores_2x2, spec=spec, hop_input=hop, ) print( "hop_output nodes:", written["node"].values.tolist(), ) print( "hop_input nodes:", read["node"].values.tolist(), ) print( "hop_output equals layer=1:", written.identical(named_hidden), ) print( "layer coord on hop_input:", "layer" in read.coords, )

hop_output nodes: ['H1', 'H2']
hop_input nodes: ['A', 'B']
hop_output equals layer=1: True
layer coord on hop_input: False

Tensor layouts¶

The axis count decides how a tensor is read.

Default 2-D (observation, node)¶

A 2-D tensor is read as (observation, node). Dim 1 must equal spec.layer_dims[layer], the unit count of that layer. The DataArray keeps the raw values; nothing is reduced.

In [ ]:

hidden = kpnn2.map_node_attributions(
    attributions=scores_2x2,
    spec=spec,
    layer=1,
)
print(
    "dims:",
    hidden.dims,
)
print(hidden)

hidden = kpnn2.map_node_attributions( attributions=scores_2x2, spec=spec, layer=1, ) print( "dims:", hidden.dims, ) print(hidden)

dims: ('observation', 'node')
<xarray.DataArray (observation: 2, node: 2)> Size: 16B
array([[0.1, 0.2],
       [0.3, 0.4]], dtype=float32)
Coordinates:
  * observation  (observation) int64 16B 0 1
  * node         (node) <U2 16B 'H1' 'H2'
    layer        int64 8B 1

1-D¶

A vector is labeled as dim (node,). That is the layout for a single observation, or for scores you already averaged yourself.

In [ ]:

one_row = torch.tensor([0.90, 1.10])
vec = kpnn2.map_node_attributions(
    attributions=one_row,
    spec=spec,
    layer=1,
)
print(
    "dims:",
    vec.dims,
)
print(vec.to_dataframe(name="score").reset_index())

one_row = torch.tensor([0.90, 1.10]) vec = kpnn2.map_node_attributions( attributions=one_row, spec=spec, layer=1, ) print( "dims:", vec.dims, ) print(vec.to_dataframe(name="score").reset_index())

dims: ('node',)
  node  layer  score
0   H1      1    0.9
1   H2      1    1.1

Extra axes (dims / coords)¶

Rank 3 or higher has no default names: nothing tells the function what the extra axis means. dims must contain node exactly once. coords may label the other axes; it must not include node or layer (layer is always the scalar you passed in).

In [ ]:

# Shape (observation, class, node) at the hidden layer.
scores_cls = torch.tensor(
    [
        [[0.10, 0.20], [0.30, 0.40]],
        [[0.50, 0.60], [0.70, 0.80]],
    ]
)
by_class = kpnn2.map_node_attributions(
    attributions=scores_cls,
    spec=spec,
    layer=1,
    dims=("observation", "class", "node"),
    coords={"class": ["neg", "pos"]},
)
print(
    "dims:",
    by_class.dims,
)
print(by_class.to_dataframe(name="score").reset_index())

# Shape (observation, class, node) at the hidden layer. scores_cls = torch.tensor( [ [[0.10, 0.20], [0.30, 0.40]], [[0.50, 0.60], [0.70, 0.80]], ] ) by_class = kpnn2.map_node_attributions( attributions=scores_cls, spec=spec, layer=1, dims=("observation", "class", "node"), coords={"class": ["neg", "pos"]}, ) print( "dims:", by_class.dims, ) print(by_class.to_dataframe(name="score").reset_index())

dims: ('observation', 'class', 'node')
   observation class node  layer  score
0            0   neg   H1      1    0.1
1            0   neg   H2      1    0.2
2            0   pos   H1      1    0.3
3            0   pos   H2      1    0.4
4            1   neg   H1      1    0.5
5            1   neg   H2      1    0.6
6            1   pos   H1      1    0.7
7            1   pos   H2      1    0.8

Several calls become step¶

Scores from several forward or Captum calls go in together. A tuple or list of equal-shaped tensors is stacked on a new step axis, one entry per call, and stacked 2-D pieces default to (step, observation, node).

In [ ]:

step0 = torch.tensor(
    [
        [0.10, 0.20],
        [0.30, 0.40],
    ]
)
step1 = torch.tensor(
    [
        [0.50, 0.60],
        [0.70, 0.80],
    ]
)
by_step = kpnn2.map_node_attributions(
    attributions=(step0, step1),
    spec=spec,
    layer=1,
)
print(
    "dims:",
    by_step.dims,
)
print(by_step.to_dataframe(name="score").reset_index())

step0 = torch.tensor( [ [0.10, 0.20], [0.30, 0.40], ] ) step1 = torch.tensor( [ [0.50, 0.60], [0.70, 0.80], ] ) by_step = kpnn2.map_node_attributions( attributions=(step0, step1), spec=spec, layer=1, ) print( "dims:", by_step.dims, ) print(by_step.to_dataframe(name="score").reset_index())

dims: ('step', 'observation', 'node')
   step  observation node  layer  score
0     0            0   H1      1    0.1
1     0            0   H2      1    0.2
2     0            1   H1      1    0.3
3     0            1   H2      1    0.4
4     1            0   H1      1    0.5
5     1            0   H2      1    0.6
6     1            1   H1      1    0.7
7     1            1   H2      1    0.8

Wrong width¶

A wrong node-axis length is an error, not a silent relabeling. Passing hidden-layer scores as layer=2 fails: C is one unit, the tensor has two.

In [ ]:

try:
    kpnn2.map_node_attributions(
        attributions=scores_2x2,
        spec=spec,
        layer=2,
    )
except kpnn2.Kpnn2Error as exc:
    print(exc)

try: kpnn2.map_node_attributions( attributions=scores_2x2, spec=spec, layer=2, ) except kpnn2.Kpnn2Error as exc: print(exc)

Attribution tensor has the wrong number of units. Expected 1, got 2.

What to map¶

Tensor you have	Argument
Input scores (layer_dims[0])	layer=0
What hop i's module returns (Captum's default side)	hop_output=spec.hops[i], same as layer=i + 1
What hop i's module reads (attribute_to_layer_input=True)	hop_input=spec.hops[i]
Output scores (layer_dims[-1])	layer=len(spec.layer_nodes) - 1
BatchNorm, dropout, unnamed module	do not map

Matching width is not proof that the scores came from those nodes. Only pass tensors whose units are spec nodes.

For Captum on a trained feedforward net, see Feedforward example Step 6. For the adjacency layout, see the Cyclic graph example, the Time-series example, and the Transformer example.

Aggregating node scores¶

map_node_attributions() keeps one value per observation. aggregate_node_attributions() folds those scores with a registered method. list_aggregation_methods() lists the registry. The dispatcher is on the API page; method pages are linked from that page.

The default method is rauter_mangano_2026. It is registered and not yet implemented. Calling it raises Kpnn2Error.

In [ ]:

kpnn2.list_aggregation_methods()[["name", "status", "description"]]

kpnn2.list_aggregation_methods()[["name", "status", "description"]]

Out[ ]:

	name	status	description
0	rauter_mangano_2026	recommended	Not yet implemented.