Skip to content

scatter_hop_outputs

scatter_hop_outputs(
    tensor: object,
    hop: Hop,
) -> dict[int, torch.Tensor]

Split a tensor laid out like one hop's input back into its layers.

gather_hop_inputs lays the layers a hop reads side by side; this function cuts such a tensor apart again. You need it when a model runs a hop backwards, as the decoder of a tied autoencoder does. PackedLinear.transpose gives the layer built from a hop its reverse map, from the hop's target layer back onto the hop's input. On a hop that reads several layers, that output is those layers glued together, and the decoder can continue layer by layer only once each piece is filed under its own depth. This function makes the cut and returns the pieces keyed by depth. Like the gather, it does no arithmetic.

Take A and B feeding H, H feeding C, and the skip A -> C. The hop into C reads layer 0 [A, B] and layer 1 [H], so its transposed layer maps C to the three columns [A, B, H], and this splits them into {0: [A, B], 1: [H]}. Going forward, layer 0 feeds two hops, the one into H and the one into C. Going backward, it therefore receives a piece from each, and the transpose of a layer read twice is the sum of the two. That is why this returns pieces instead of writing them anywhere: the adding stays in your code.

Precisely, this is the inverse of gather_hop_inputs for that axis. The encoder concatenates whole source layers; a tied decoder (PackedLinear.transpose) emits the same concatenated width. This splits it. It does not take a saved dict and does not add into one: return the pieces and add them yourself, because two reversed hops may write the same earlier layer.

Parameters:

Name Type Description Default

tensor

Tensor

Concatenated source axis, last dimension hop.in_features. Typically the output of PackedLinear.transpose() on this hop. Leading dimensions are the caller's, typically a batch.

required

hop

Hop

The hop whose source axis tensor follows, one entry of spec.hops.

required

Returns:

Type Description
dict of int to torch.Tensor

One entry per hop.source_layers, in that order. result[layer] has last dimension hop.source_dims[i] for that layer. A hop with one source layer returns tensor itself rather than a copy, so writing into the piece writes into tensor. Several source layers are torch.split views on the last axis.

Raises:

Type Description
Kpnn2Error

If hop is not a Hop; tensor is not a tensor; tensor is 0-dimensional; or the last dimension is not hop.in_features.

See Also

gather_hop_inputs : Concatenates the layers this splits. PackedLinear.transpose : Packed W.T; its output on a skip hop is what this splits. Hop : source_layers, source_dims, and column_offsets are the split layout. PackedLinear : Encoder map whose transpose emits this axis.

Notes

Nothing here is graph-aware and nothing holds weights. Unused skip columns stay in the pieces, same as they stay in a gather. An AdjacencySpec has no hops and is not accepted.

Examples:

The hop into C reads layers 0 and 1, so a concatenated decoder activation splits back onto those depths:

>>> import pandas as pd
>>> import torch
>>> import kpnn2
>>> edgelist = pd.DataFrame(
...     {
...         "source": ["A", "H", "A"],
...         "target": ["H", "C", "C"],
...     }
... )
>>> spec = kpnn2.parse_layered(edgelist)
>>> concat = torch.tensor([[2.0, 5.0]])
>>> parts = kpnn2.scatter_hop_outputs(
...     concat,
...     spec.hops[1],
... )
>>> list(parts)
[0, 1]
>>> parts[0].tolist(), parts[1].tolist()
([[2.0]], [[5.0]])