Skip to content

gather_hop_inputs

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

Build one hop's input from the layer activations you saved.

Without skip edges, each layer reads only the layer directly below it, so one module's output is the next module's input. A skip edge breaks that chain: its source sits two or more layers below its target. kpnn2 still applies every edge into a layer with one module, the one built from that layer's hop (PackedLinear.from_hop(hop) or MaskedLinear(hop.to_mask())). That module therefore needs every layer its target's parents sit in, at once and as one tensor. This function builds that tensor. It does no arithmetic; it only arranges tensors you already have.

Take A and B feeding H, H feeding C, and the skip A -> C. Layer 0 is [A, B], layer 1 is [H], and the hop into C reads both, so its input is the three columns [A, B, H]: whole layers, laid side by side in depth order. B is among them although there is no edge B -> C. The hop says which layers to read, in which order, and how many units wide each must be. Without skip edges, every hop reads a single layer, and that layer is the input as it stands.

Precisely, a hop's source columns are whole source layers laid side by side, so the tensor feeding PackedLinear or MaskedLinear on that hop is those layers concatenated. Call it between hops, keeping every layer you produce in saved; a forgotten layer raises rather than dropping the edges that read it. Inputs are never modified, and a single-source hop returns the saved tensor itself.

Parameters:

Name Type Description Default

saved

mapping of int to torch.Tensor

Layer depth to that layer's activation, as far as forward() has produced them. saved[i] is shaped (..., layer_dims[i]); the leading dimensions are the caller's, typically a batch. Only the layers in hop.source_layers are read, so extra keys are ignored, and those layers must share a dtype, a device, and their leading dimensions. Neither the mapping nor its tensors are copied or modified.

required

hop

Hop

The hop about to be applied, one entry of spec.hops. Its source_layers and source_dims decide which keys are read, in which order, and how wide each one must be; the packed indices are not used here.

required

Returns:

Type Description
Tensor

The hop's source axis, shape (..., hop.in_features), columns in concatenated source-unit order, dtype and device of the saved layers. source_nodes is one name per node, so it is shorter than this axis when a source node is wider than 1. When the hop reads a single layer, including hops[0] and any hop whose parents all sit at one depth, this is that saved tensor itself rather than a copy, so writing into it writes into the saved activation.

Raises:

Type Description
Kpnn2Error

If saved is not a mapping or hop is not a Hop; a layer in hop.source_layers is absent from saved or is not a tensor; a saved tensor is 0-dimensional or is the wrong number of units wide; or the source layers disagree on dtype, device, or a dimension other than the last, such as the batch size.

See Also

scatter_hop_outputs : Split the concatenated axis this returns back onto source layers. PackedLinear : Applies the hop to the tensor returned here. PackedLinear.transpose : Tied decode of that hop; feed its output to scatter_hop_outputs when the hop reads several layers. MaskedLinear : Dense hatch via hop.to_mask(). Hop : The record that fixes the source layers and the concatenated column order this follows. align_inputs : Column index that puts named features in the order layer 0 expects.

Notes

Nothing here is graph-aware: this holds no weights, does not inject values into the previous layer, and does not pick skip sources by name. Columns that are not edges stay in the concatenated tensor; PackedLinear never reads them, and MaskedLinear(hop.to_mask()) zeros them. In the [A, B, H] example above, the A -> C weight is the packed pair whose source column is A.

An AdjacencySpec has no hops and is not accepted; in that layout every edge is already a packed index pair.

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. ...