Mapping attributions¶

map_node_attributions() labels a tensor with spec node names. The result is 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 Getting started Step 6. For an AdjacencySpec, see the Recurrent example. The signature is on the API page.

The naming rule¶

You compute the tensor (Captum LayerConductance, input gradients, or any other scores with the same width as that layer). This function only names the node axis.

Captum scores modules. LayeredSpec scores layers. layer= is an index into spec.layer_nodes, not a Captum target name. Layer 0 is the input nodes. A MaskedLinear built from spec.hops[i].mask writes the activations at layer hops[i].target_layer, that is i + 1. Pass that for the hop's output scores. A hop may read several layers, but it writes exactly one.

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

With an AdjacencySpec you omit layer=: that layout has no depths, so the node axis is the whole state vector (spec.nodes) and the result carries no layer coordinate. See the Recurrent example.

The graph below makes layer= mandatory. 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 can share a width. The function does not guess the layer from tensor.shape[-1].

In [1]:

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

spec.hops[i].mask is the incoming hop that writes layer i + 1. The hop 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 [2]:

hop0 = kpnn2.MaskedLinear(spec.hops[0].mask)
hop1 = kpnn2.MaskedLinear(spec.hops[1].mask)
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.MaskedLinear(spec.hops[0].mask) hop1 = kpnn2.MaskedLinear(spec.hops[1].mask) 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 [3]:

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

Tensor layouts¶

Default 2-D (observation, node)¶

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

In [4]:

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 [5]:

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. 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 [6]:

# 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¶

A tuple or list of equal-shaped tensors is stacked on a new step axis (one entry per forward / Captum call). Default names for stacked 2-D pieces are (step, observation, node).

In [7]:

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¶

Passing hidden-layer scores as layer=2 fails: C is one unit, the tensor has two.

In [8]:

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	layer=
Input scores (layer_dims[0])	0
Output of MaskedLinear(spec.hops[i].mask)	i + 1
Output scores (layer_dims[-1])	last layer
BatchNorm, dropout, unnamed module	do not map

Last layer is len(spec.layer_nodes) - 1. 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 Getting started Step 6. For the adjacency layout, see the Recurrent example.