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 (a wide node's name repeats once per unit).	required
layer	int	0-based depth into spec.layer_nodes, index 0 being the input layer, whose names label the node axis; the index itself is attached as a scalar layer coordinate. 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. No layer coordinate is attached; that axis is not one depth. 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, scalar 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 scalar 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. 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. 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), and a scalar layer coordinate when layer or hop_output was passed, and when axis="inputs" is passed on a LayeredSpec. A hop_input mapping, and axis="inputs" on an AdjacencySpec, have 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.

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 int 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, 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 names an unknown axis or a 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]
>>> int(da.coords["layer"])
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']
>>> int(out_da.coords["layer"])
1
>>> in_da = kpnn2.map_node_attributions(
...     attributions=hop_scores,
...     spec=square_spec,
...     hop_input=hop,
... )
>>> in_da["node"].values.tolist()
['A', 'B']
>>> "layer" in in_da.coords
False

A skip hop reads several layers, so its input axis is those layers concatenated (width hop.in_features):

>>> 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']

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