map_node_attributions

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

Label an attribution tensor's node axis with names from a spec.

Attribution methods return unlabeled tensors whose node axis is bare positions; the spec knows the name at each one. Reach for it after Captum or your own gradients rather than zipping names to columns yourself. Which spec you pass decides the contract: a LayeredSpec needs exactly one of layer, hop_input, hop_output, or axis="inputs"; an AdjacencySpec forbids the first three. Omit axis there to name the whole state vector, or pass axis="inputs" to name the input units. The width is never inferred. Values are detached onto CPU and never aggregated.

Parameters:

Name	Type	Description	Default
attributions	torch.Tensor or sequence of torch.Tensor	Scores exactly as the attribution method produced them, in whatever units it works in; nothing is scaled, summed, or made absolute. The node axis must be as long as the named units — spec.layer_dims[layer] when layer is given, hop.in_features for hop_input=hop, hop.out_features for hop_output=hop, len(spec.input_index) for axis="inputs" on an AdjacencySpec, spec.layer_dims[0] for axis="inputs" on a LayeredSpec (the same axis as layer=0), and spec.state_dim for an AdjacencySpec when axis is omitted — and the remaining axes are yours. A non-empty tuple or list of equal-shaped tensors is stacked on a new leading step axis, one entry per unrolled step or module call. The tensors are read, never modified, and the result holds a detached CPU copy that shares no memory with them.	required
spec	LayeredSpec or AdjacencySpec	Parsed edgelist supplying the node coordinate. A LayeredSpec names one depth (layer or hop_output) or the concatenated source units of one hop (hop_input); an AdjacencySpec names the whole state vector at once and has no depth to report, unless axis="inputs", which names spec.input_nodes, one unit each.	required
layer	int or numpy integer	0-based depth into spec.layer_nodes, index 0 being the input layer, whose names label the node axis; the layer coordinate repeats the index on every unit, the same for a numpy integer as for an int. LayeredSpec only, and mutually exclusive with hop_input and hop_output.	None
hop_input	Hop	A Hop that equals one entry of spec.hops. Labels what that hop's module reads: the concatenated source axis gather_hop_inputs returns, length hop.in_features, with unit names (a wide node's name repeats). Captum layer methods score this side only with attribute_to_layer_input=True. The layer coordinate gives each unit the depth of its source layer, so a skip hop's axis holds several depths. LayeredSpec only, and mutually exclusive with layer and hop_output.	None
hop_output	Hop	A Hop that equals one entry of spec.hops. Labels what that hop's module returns: layer hop.target_layer, length hop.out_features. This is the side Captum layer methods score by default. The result is the same as layer=hop.target_layer, layer coordinate included. LayeredSpec only, and mutually exclusive with layer and hop_input.	None
axis	inputs	"inputs" labels the model-input axis and nothing else. On a LayeredSpec that is layer=0, including the layer coordinate, and it is mutually exclusive with layer, hop_input, and hop_output. On an AdjacencySpec the names are the input units in spec.input_index order, length len(spec.input_index), with no layer coordinate. Any other string is rejected. The length of the tensor is not used to choose this axis.	"inputs"
dims	sequence of str	One name per axis of the tensor after any stacking, containing node exactly once and never layer. A single str is rejected, not split into characters, and so are a set and a mapping, which carry no axis order. Required at 3 or more axes, unless the tensor is a stacked sequence of 1-D or 2-D pieces. The defaults are ("node",) for 1-D and ("observation", "node") for 2-D, with step prepended when a sequence was stacked.	None
coords	mapping of str to sequence	Labels for axes other than node and layer, keyed by dim name; each sequence must be as long as its axis. A single str is rejected, not split into characters, and so are a set and a mapping (a dict would give its keys). Axes left out are labelled with their integer positions.	None

Returns:

Type

Description

DataArray

The values and shape of the (stacked) tensor, carrying the spec's node names as the node coordinate in spec order (a name repeats once per unit when that node is wider than 1). On a LayeredSpec it also carries layer, an int64 coordinate along node that holds each unit's depth in spec.layer_nodes: one depth repeated for layer, hop_output, and axis="inputs", and each unit's source layer for hop_input. It stays with its units when you concatenate results along node. Read one depth with int(da["layer"][0]). An AdjacencySpec result has no layer coordinate. bfloat16 scores are stored as float32: NumPy has no bfloat16 dtype, and every bfloat16 value fits in float32, so the scores are unchanged. Other dtypes are kept. Use .to_dataframe(name="score").reset_index() for a long table, or .to_pandas() for a 2-D wide table. Units of a wide node share its edges, so the spec gives them no names of their own. On an axis with a wide node the node index is not unique. da.sel(node=name) returns all units of that node, and a fold you choose, such as da.groupby("node").mean(), gives one score per node, sorted by name. Selecting a list of names, reindex, and aligning with an array whose node index differs raise in xarray, even for width-1 names. Select several names with da.sel(node=da["node"].isin(names)). Two arrays with the same node index combine unit by unit, by position. With every node at width 1 the index is unique.

Raises:

Type

Description

Kpnn2Error

If spec is neither a LayeredSpec nor an AdjacencySpec; a LayeredSpec gets none or more than one of layer, hop_input, hop_output, and axis; axis is neither "inputs" nor None; any of layer, hop_input, or hop_output is given for an AdjacencySpec; layer is not an integer in range; hop_input or hop_output is not a Hop that matches an entry of spec.hops; attributions is neither a tensor nor a non-empty sequence of equal-shaped tensors; dims is missing, not a sequence of strings (a str, set, or mapping included), the wrong length, non-unique, or does not name node exactly once; the node axis is not as long as the named units; or coords is not a mapping, names an unknown axis, or holds a value that is not a sequence of labels (a str, set, or mapping included) or has the wrong length.

See Also

aggregate_node_attributions : Folds named scores to one value per node; call this after mapping. align_inputs : Shared input_nodes order on the way in; a column index into the caller's feature names, not the full state axis of an AdjacencySpec. gather_hop_inputs : Builds the tensor whose columns hop_input names. Hop : One hop; pass hop_output=spec.hops[i] or hop_input=spec.hops[i] to name either side of it. LayeredSpec : Holds layer_nodes and hops, the names used when layer, hop_input, or hop_output is given. AdjacencySpec : Holds nodes, the state-vector names used when all three are omitted.

Notes

Captum is not imported anywhere in this package; mapping is name alignment only, and any attribution method will do.

Name a hop module's scores with the same index i you use for the module. For what PackedLinear (or MaskedLinear) on spec.hops[i] returns, which is what Captum layer methods such as LayerConductance score by default, pass hop_output=spec.hops[i]. For what it reads (the gather_hop_inputs tensor, or Captum with attribute_to_layer_input=True), pass hop_input=spec.hops[i]. The side is always stated, never inferred: a hop's input and output often have the same width, so a width check alone cannot catch the wrong side. Only map units that are spec nodes; BatchNorm and other unnamed modules have no node axis to name.

DeepLift, DeepLiftShap, and LRP rescale nn.Module nonlinearities. The activation after a hop is an nn.ReLU held in an nn.ModuleList, one entry per hop, and called from forward. The hop module returns the pre-activation tensor: hop_output names that linear map. Post-activation node states are the output of the nn.ReLU that follows the hop. LayerActivation and LayerGradientXActivation hook that module.

A recurrent net on an AdjacencySpec has no layer to index, and the natural extra axis there is step: pass one tensor per unrolled step as a sequence and they are stacked for you.

Examples:

Name a two-observation tensor at the output layer of a LayeredSpec:

>>> 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]
>>> da["layer"].values.tolist()
[2]

Both sides of a hop can have the same width. Here hops[0] reads A, B and writes H1, H2, so one (1, 2) tensor fits either side; the keyword decides the names:

>>> square = pd.DataFrame(
...     {
...         "source": ["A", "B", "A", "B", "H1", "H2"],
...         "target": ["H1", "H1", "H2", "H2", "C", "C"],
...     }
... )
>>> square_spec = kpnn2.parse_layered(square)
>>> hop = square_spec.hops[0]
>>> hop.in_features, hop.out_features
(2, 2)
>>> hop_scores = torch.tensor([[0.1, 0.2]])
>>> out_da = kpnn2.map_node_attributions(
...     attributions=hop_scores,
...     spec=square_spec,
...     hop_output=hop,
... )
>>> out_da["node"].values.tolist()
['H1', 'H2']
>>> out_da["layer"].values.tolist()
[1, 1]
>>> in_da = kpnn2.map_node_attributions(
...     attributions=hop_scores,
...     spec=square_spec,
...     hop_input=hop,
... )
>>> in_da["node"].values.tolist()
['A', 'B']
>>> in_da["layer"].values.tolist()
[0, 0]

A skip hop reads several layers, so its input axis is those layers concatenated (width hop.in_features), and layer gives each unit its own depth:

>>> skip_edges = pd.DataFrame(
...     {
...         "source": ["A", "H", "A"],
...         "target": ["H", "C", "C"],
...     }
... )
>>> skip_spec = kpnn2.parse_layered(skip_edges)
>>> skip_hop = skip_spec.hops[1]
>>> skip_hop.source_layers
(0, 1)
>>> skip_da = kpnn2.map_node_attributions(
...     attributions=torch.tensor([[0.1, 0.2]]),
...     spec=skip_spec,
...     hop_input=skip_hop,
... )
>>> skip_da["node"].values.tolist()
['A', 'H']
>>> skip_da["layer"].values.tolist()
[0, 1]

Concatenated along node, each unit keeps its depth:

>>> import xarray as xr
>>> top_da = kpnn2.map_node_attributions(
...     attributions=torch.tensor([[0.3]]),
...     spec=skip_spec,
...     layer=2,
... )
>>> both = xr.concat(
...     [skip_da, top_da],
...     dim="node",
... )
>>> both["node"].values.tolist()
['A', 'H', 'C']
>>> both["layer"].values.tolist()
[0, 1, 2]

A wide node's name repeats once per unit. sel takes one name; pick several with isin, and fold units yourself:

>>> fan = pd.DataFrame(
...     {
...         "source": ["A", "A", "A", "H1", "H2", "H3"],
...         "target": ["H1", "H2", "H3", "C", "C", "C"],
...     }
... )
>>> fan_spec = kpnn2.parse_layered(
...     fan,
...     widths={"H1": 2},
... )
>>> fan_da = kpnn2.map_node_attributions(
...     attributions=torch.tensor([[1.0, 2.0, 3.0, 4.0]]),
...     spec=fan_spec,
...     layer=1,
... )
>>> fan_da["node"].values.tolist()
['H1', 'H1', 'H2', 'H3']
>>> fan_da.sel(node="H1").values.tolist()
[[1.0, 2.0]]
>>> picked = fan_da.sel(node=fan_da["node"].isin(["H1", "H3"]))
>>> picked["node"].values.tolist()
['H1', 'H1', 'H3']
>>> fan_da.groupby("node").sum().values.tolist()
[[3.0, 3.0, 4.0]]

On an AdjacencySpec there are no layers: omit layer, hop_input, and hop_output and the whole state vector is named. One tensor per unrolled 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
>>> input_scores = kpnn2.map_node_attributions(
...     attributions=torch.tensor([[0.5], [1.0]]),
...     spec=state_spec,
...     axis="inputs",
... )
>>> input_scores["node"].values.tolist()
['x']
>>> "layer" in input_scores.coords
False