Supported architectures¶
kpnn2 supplies connectivity primitives, not finished models: you
write an ordinary nn.Module and call them from your own
forward(). The table below says which architecture families you
can build that way.
| Type | Supported | What to use |
|---|---|---|
| Feedforward | Yes | parse_layered (optional widths= to give a node several units), PackedLinear, gather_hop_inputs; tied decode PackedLinear.transpose + scatter_hop_outputs → Feedforward example |
| Cyclic | Yes | parse_adjacency, PackedLinear or MaskedLinear(spec.to_mask()), loop you own → Cyclic graph example |
| Transformer | Yes | parse_adjacency + PackedMultiheadAttention; encoder / FFN / head are yours → Transformer example |
| Sequence RNN / GRU / LSTM | Yes, see Sequence models | parse_adjacency + MaskedLinear / PackedLinear as the maps; you write the time loop (not nn.RNN / nn.GRU / nn.LSTM) → Time-series example |
| Graph NN | No | Knowledge-primed GNNs are a natural edgelist model in PyG; we will not build or maintain kpnn-pyg. |
| Convolutional NN | No | No sparsely connected support. Build it yourself in PyTorch (nn.Conv1d / nn.Conv2d / nn.Conv3d) |
Sequence models¶
The sequence row needs a longer answer: PyTorch's recurrent
modules cannot carry a prior. nn.RNN, nn.GRU, and nn.LSTM
are fused convenience modules whose maps are dense W_ih /
W_hh. None accepts an edgelist, so none can be the recurrent map
of a knowledge-primed neural network (KPNN);
this package does not ship MaskedRNN, MaskedGRU, or
MaskedLSTM.
What replaces them is the cyclic-graph recipe with a changing
x_t at each step. Memory lives in the edgelist; the time loop is
yours:
- Parse with
parse_adjacency(). A self-loop or hidden→hidden edge lets a named node carry state across time;parse_layered()is DAG-only and rejects both. - One map, shared across steps:
MaskedLinear(spec.to_mask()), orPackedLinearwhennis large, sinceto_mask()allocates(n, n). - In your
forward(), loop over timet. Registertorch.as_tensor(spec.input_index)as a buffer and writex_twithindex_copyinto the state vector, then apply the map. Each step Captum should hook is annn.Identity. The shared map stays a normal child of the module. - You choose
n_steps, the sequence length.kpnn2does not unroll time or pick a step count.
An Elman-style net is that loop. A GRU or LSTM is the usual gate
equations with one MaskedLinear per map; whether every gate
shares the same prior is your modeling choice.
state = state.index_copy(
-1,
self.input_index,
x_t,
)
state = tap(self.act(self.core(state)))
tap is one nn.Identity from an nn.ModuleList. The
Time-series example walks through
the Elman loop, attributions on a
step axis, and a no-memory control.