spikeforge 0.3.0__tar.gz → 0.3.2__tar.gz

{spikeforge-0.3.0 → spikeforge-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: spikeforge
-Version: 0.3.0
+Version: 0.3.2
 Summary: Spiking neural network trainer/experiments for MNIST built with snnTorch and PyTorch.
 Author-email: Capsize LLC <contact@capsizegames.com>
 Maintainer-email: Capsize LLC <contact@capsizegames.com>
@@ -56,6 +56,7 @@ Dynamic: license-file
 # spikeforge
 
 [![CI](https://github.com/capsize-games/spikeforge/actions/workflows/ci.yml/badge.svg)](https://github.com/capsize-games/spikeforge/actions/workflows/ci.yml)
+[![Discord](https://img.shields.io/badge/Discord-Join%20the%20community-5865F2?logo=discord&logoColor=white)](https://capsizegames.com/discord)
 [![Status: pre-1.0](https://img.shields.io/badge/status-pre--1.0-orange.svg)](OPEN_SOURCE_CHECKLIST.md)
 [![License: BSD-3-Clause](https://img.shields.io/badge/License-BSD_3--Clause-blue.svg)](LICENSE)
 [![Python 3.10–3.13](https://img.shields.io/badge/python-3.10%20%7C%203.11%20%7C%203.12%20%7C%203.13-3776AB.svg)](https://www.python.org/downloads/)
@@ -63,15 +64,15 @@ Dynamic: license-file
 [![PRs Welcome](https://img.shields.io/badge/PRs-welcome-brightgreen.svg)](https://github.com/capsize-games/spikeforge/blob/main/CONTRIBUTING.md)
 [![Docs](https://img.shields.io/badge/docs-long--form%20reference-blue.svg)](documentation/README.md)
 
+![spikeforge dashboard](images/dashboard.png)
+
 A spiking-neural-network (SNN) toolkit built on
 [snnTorch](https://snntorch.readthedocs.io/) and PyTorch. Loads MNIST-style
 and neuromorphic event datasets, encodes them into rate, latency, delta, and
 random spikes, and trains, validates, exports, and deploys LIF networks —
 with a live browser dashboard served over WebSockets.
 
-> **Pre-1.0 and unpublished.** `install.sh` is the supported path today;
-> PyPI ships once the distributions are published. Before trusting any
-> number this produces, read
+> **Pre-1.0.** Before trusting any number this produces, read
 > [Implications and boundaries](documentation/implications-and-boundaries.md).
 
 ## Quickstart

{spikeforge-0.3.0 → spikeforge-0.3.2}/README.md

@@ -1,6 +1,7 @@
 # spikeforge
 
 [![CI](https://github.com/capsize-games/spikeforge/actions/workflows/ci.yml/badge.svg)](https://github.com/capsize-games/spikeforge/actions/workflows/ci.yml)
+[![Discord](https://img.shields.io/badge/Discord-Join%20the%20community-5865F2?logo=discord&logoColor=white)](https://capsizegames.com/discord)
 [![Status: pre-1.0](https://img.shields.io/badge/status-pre--1.0-orange.svg)](OPEN_SOURCE_CHECKLIST.md)
 [![License: BSD-3-Clause](https://img.shields.io/badge/License-BSD_3--Clause-blue.svg)](LICENSE)
 [![Python 3.10–3.13](https://img.shields.io/badge/python-3.10%20%7C%203.11%20%7C%203.12%20%7C%203.13-3776AB.svg)](https://www.python.org/downloads/)
@@ -8,15 +9,15 @@
 [![PRs Welcome](https://img.shields.io/badge/PRs-welcome-brightgreen.svg)](https://github.com/capsize-games/spikeforge/blob/main/CONTRIBUTING.md)
 [![Docs](https://img.shields.io/badge/docs-long--form%20reference-blue.svg)](documentation/README.md)
 
+![spikeforge dashboard](images/dashboard.png)
+
 A spiking-neural-network (SNN) toolkit built on
 [snnTorch](https://snntorch.readthedocs.io/) and PyTorch. Loads MNIST-style
 and neuromorphic event datasets, encodes them into rate, latency, delta, and
 random spikes, and trains, validates, exports, and deploys LIF networks —
 with a live browser dashboard served over WebSockets.
 
-> **Pre-1.0 and unpublished.** `install.sh` is the supported path today;
-> PyPI ships once the distributions are published. Before trusting any
-> number this produces, read
+> **Pre-1.0.** Before trusting any number this produces, read
 > [Implications and boundaries](documentation/implications-and-boundaries.md).
 
 ## Quickstart

{spikeforge-0.3.0 → spikeforge-0.3.2}/pyproject.toml

@@ -15,7 +15,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "spikeforge"
-version = "0.3.0"
+version = "0.3.2"
 description = "Spiking neural network trainer/experiments for MNIST built with snnTorch and PyTorch."
 # ``README.md`` and ``LICENSE`` are committed symlinks to the repository-root
 # files: the distribution metadata points at the root documents while

spikeforge-0.3.2/spikeforge/memory/embedding_readout.py

@@ -0,0 +1,51 @@
+"""Bridge a spiking embedder's output spikes into memory-module calls.
+
+Unlike :mod:`spikeforge.memory.lockstep_runner` (issue #22, a frozen
+*classifier* with a hidden layer feeding memory), an episodic embedder
+(issue #24) has no separate classification head: its own output stage
+*is* the spike pattern memory reads.
+"""
+
+from typing import Tuple
+
+import torch
+
+from spikeforge.memory.one_shot_associative_memory import (
+    OneShotAssociativeMemory,
+)
+from spikeforge.simulator.module_spec import spec_of
+from spikeforge.simulator.runner import run
+from spikeforge.simulator.trajectory import Trajectory
+from spikeforge.topology.stage_module import StageModule
+
+
+def output_spike_trace(
+    net: StageModule, spikes: torch.Tensor,
+) -> Tuple[torch.Tensor, Trajectory]:
+    """Return ``([T, B, D]`` output spikes, full trajectory)`` for ``net``."""
+    trajectory = run(net, spikes, track=True)
+    return trajectory.spikes[spec_of(net).output], trajectory
+
+
+def teach_from_example(
+    memory: OneShotAssociativeMemory,
+    net: StageModule,
+    spikes: torch.Tensor,
+) -> int:
+    """One-shot teach ``memory`` from one example's encoded ``spikes``."""
+    trace, _ = output_spike_trace(net, spikes)
+    return memory.teach(trace[:, 0, :])
+
+
+def memory_predictions(
+    memory: OneShotAssociativeMemory,
+    net: StageModule,
+    spikes: torch.Tensor,
+) -> torch.Tensor:
+    """Return ``[B]`` predicted class indices from memory alone."""
+    trace, _ = output_spike_trace(net, spikes)
+    memory.reset_state(trace.size(1))
+    counts = torch.zeros(trace.size(1), memory.num_classes)
+    for step in trace:
+        counts += memory.step(step)
+    return counts.argmax(dim=1)

spikeforge-0.3.2/spikeforge/memory/episode_sampler.py

@@ -0,0 +1,63 @@
+"""N-way K-shot episode sampling for episodic metric-learning (issue #24)."""
+
+import random
+from typing import Dict, List, Sequence, Tuple
+
+import torch
+from torch.utils.data import Dataset
+
+#: One sampled episode: support images/labels, then query images/labels.
+Episode = Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
+
+
+def class_indices(
+    dataset: Dataset, classes: Sequence[int],
+) -> Dict[int, List[int]]:
+    """Return each of ``classes``'s sample indices within ``dataset``."""
+    targets = dataset.targets
+    return {
+        cls: torch.nonzero(targets == cls).squeeze(1).tolist()
+        for cls in classes
+    }
+
+
+def sample_episode(
+    dataset: Dataset,
+    by_class: Dict[int, List[int]],
+    n_way: int,
+    k_shot: int,
+    n_query: int,
+    rng: random.Random,
+) -> Episode:
+    """Sample one episode: ``n_way`` classes, ``k_shot`` + ``n_query`` each.
+
+    Labels are episode-local (``0..n_way-1``), independent of the
+    dataset's real class ids -- the point of episodic training is a
+    task-agnostic "these match / these don't" signal, not a fixed
+    N-way classifier.
+    """
+    chosen = rng.sample(list(by_class), n_way)
+    support: Tuple[List[torch.Tensor], List[int]] = ([], [])
+    query: Tuple[List[torch.Tensor], List[int]] = ([], [])
+    for label, cls in enumerate(chosen):
+        indices = rng.sample(by_class[cls], k_shot + n_query)
+        _fill(dataset, indices[:k_shot], label, support)
+        _fill(dataset, indices[k_shot:], label, query)
+    return (
+        torch.stack(support[0]), torch.tensor(support[1]),
+        torch.stack(query[0]), torch.tensor(query[1]),
+    )
+
+
+def _fill(
+    dataset: Dataset,
+    indices: List[int],
+    label: int,
+    accumulator: Tuple[List[torch.Tensor], List[int]],
+) -> None:
+    """Append ``dataset``'s images at ``indices`` under one episode label."""
+    images, labels = accumulator
+    for index in indices:
+        image, _ = dataset[index]
+        images.append(image)
+        labels.append(label)

spikeforge-0.3.2/spikeforge/memory/episodic_trainer.py

@@ -0,0 +1,77 @@
+"""Episodic (metric-learning) training loop for a spiking embedder."""
+
+import random
+from typing import Dict, List, Sequence
+
+import torch
+from torch.utils.data import Dataset
+
+from spikeforge.encoding.spike_encoder import SpikeEncoder
+from spikeforge.memory.episode_sampler import class_indices, sample_episode
+from spikeforge.memory.prototypical import embed, prototypical_loss
+from spikeforge.topology.registry import build_topology
+from spikeforge.topology.stage_module import StageModule
+
+#: Learning rate for the embedder's Adam optimizer.
+LEARNING_RATE = 1e-3
+
+
+def build_episodic_embedder(
+    dataset: Dataset,
+    classes: Sequence[int],
+    encoder: SpikeEncoder,
+    embed_dim: int,
+    hidden: int,
+    episodes: int,
+    n_way: int,
+    k_shot: int,
+    n_query: int,
+    seed: int = 0,
+) -> StageModule:
+    """Train ``fc_small`` as an embedder via prototypical-network episodes.
+
+    Unlike ``fc_small``'s ordinary use as a fixed-class classifier,
+    the final layer here carries no softmax semantics: episodic
+    training only ever asks it to place same-class spike patterns
+    close together and different classes apart, which is what should
+    let it generalise -- to classes held out of every episode, and in
+    the fuller claim this issue tests, to a script never trained on
+    at all (see ``few_shot_generalization_poc.py``).
+    """
+    _, net = build_topology(
+        "fc_small", {"hidden": hidden, "num_classes": embed_dim},
+    )
+    optimizer = torch.optim.Adam(net.parameters(), lr=LEARNING_RATE)
+    rng = random.Random(seed)
+    by_class = class_indices(dataset, classes)
+    for _ in range(episodes):
+        _train_one_episode(
+            net, optimizer, dataset, by_class, encoder,
+            n_way, k_shot, n_query, rng,
+        )
+    return net
+
+
+def _train_one_episode(
+    net: StageModule,
+    optimizer: torch.optim.Optimizer,
+    dataset: Dataset,
+    by_class: Dict[int, List[int]],
+    encoder: SpikeEncoder,
+    n_way: int,
+    k_shot: int,
+    n_query: int,
+    rng: random.Random,
+) -> None:
+    """Sample one episode and take one prototypical-loss gradient step."""
+    support_x, support_y, query_x, query_y = sample_episode(
+        dataset, by_class, n_way, k_shot, n_query, rng,
+    )
+    support_emb = embed(net, encoder, support_x)
+    query_emb = embed(net, encoder, query_x)
+    loss, _ = prototypical_loss(
+        support_emb, support_y, query_emb, query_y, n_way,
+    )
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()

spikeforge-0.3.2/spikeforge/memory/few_shot_generalization_poc.py

@@ -0,0 +1,148 @@
+"""Few-shot generalization to never-trained characters (issue #24).
+
+Trains a spiking embedder episodically on EMNIST letters (English
+script), freezes it, then tests whether it supports one-shot
+discrimination among KMNIST (Japanese Kuzushiji) characters -- a
+script the embedder never saw in any form during training. The
+readout is #22's unchanged
+:class:`~spikeforge.memory.one_shot_associative_memory.\
+OneShotAssociativeMemory`: one memory neuron per class in the
+evaluation episode, taught from a single support example and scored
+on held-out queries.
+
+This is a genuinely harder and uncertain claim than issue #22: no
+amount of clever memory retrieval helps distinguish two classes the
+embedder was never shown anything resembling; only a spike-pattern
+embedding that generalises across scripts can.
+"""
+
+import random
+from dataclasses import dataclass
+from typing import Dict, List
+
+import torch
+from torch.utils.data import Dataset
+
+from spikeforge.data.datasets import build_dataset
+from spikeforge.encoding.spike_encoder import SpikeEncoder
+from spikeforge.memory.embedding_readout import (
+    memory_predictions,
+    teach_from_example,
+)
+from spikeforge.memory.episode_sampler import class_indices, sample_episode
+from spikeforge.memory.episodic_trainer import build_episodic_embedder
+from spikeforge.memory.one_shot_associative_memory import (
+    OneShotAssociativeMemory,
+)
+from spikeforge.topology.stage_module import StageModule
+
+#: The script the embedder is trained on (English letters, 26 classes).
+#: torchvision's EMNIST "letters" split labels classes 1-26, not 0-25.
+TRAIN_DATASET = "emnist_letters"
+TRAIN_CLASSES = tuple(range(1, 27))
+#: The script it is evaluated on -- never seen during training.
+EVAL_DATASET = "kmnist"
+EVAL_CLASSES = tuple(range(10))
+
+
+@dataclass(frozen=True)
+class FewShotResult:
+    """Mean one-shot discrimination accuracy over held-out episodes."""
+
+    accuracy: float
+    chance: float
+    episodes: int
+
+
+def run_few_shot_generalization_poc(
+    seed: int = 0,
+    hidden: int = 64,
+    embed_dim: int = 32,
+    num_steps: int = 15,
+    train_episodes: int = 300,
+    train_n_way: int = 5,
+    train_k_shot: int = 5,
+    train_n_query: int = 5,
+    eval_episodes: int = 50,
+    eval_n_way: int = 5,
+    eval_n_query: int = 5,
+) -> FewShotResult:
+    """Train on EMNIST letters, evaluate one-shot recall on KMNIST."""
+    torch.manual_seed(seed)
+    encoder = SpikeEncoder(coding="rate", num_steps=num_steps)
+    net = _train_embedder(
+        encoder, hidden, embed_dim, train_episodes,
+        train_n_way, train_k_shot, train_n_query, seed,
+    )
+    accuracy = _evaluate_cross_script(
+        net, encoder, embed_dim, eval_episodes, eval_n_way,
+        eval_n_query, seed,
+    )
+    return FewShotResult(accuracy, 1.0 / eval_n_way, eval_episodes)
+
+
+def _train_embedder(
+    encoder: SpikeEncoder,
+    hidden: int,
+    embed_dim: int,
+    episodes: int,
+    n_way: int,
+    k_shot: int,
+    n_query: int,
+    seed: int,
+) -> StageModule:
+    """Episodically train and freeze the embedder on EMNIST letters."""
+    train_data = build_dataset(TRAIN_DATASET, train=True)
+    net = build_episodic_embedder(
+        train_data, TRAIN_CLASSES, encoder, embed_dim, hidden,
+        episodes, n_way, k_shot, n_query, seed,
+    )
+    net.requires_grad_(False)
+    net.eval()
+    return net
+
+
+def _evaluate_cross_script(
+    net: StageModule,
+    encoder: SpikeEncoder,
+    embed_dim: int,
+    eval_episodes: int,
+    n_way: int,
+    n_query: int,
+    seed: int,
+) -> float:
+    """Average one-shot discrimination accuracy over held-out episodes."""
+    eval_data = build_dataset(EVAL_DATASET, train=False)
+    by_class = class_indices(eval_data, EVAL_CLASSES)
+    rng = random.Random(seed + 1)
+    scores = [
+        _run_one_episode(
+            net, encoder, eval_data, by_class, embed_dim, n_way,
+            n_query, rng,
+        )
+        for _ in range(eval_episodes)
+    ]
+    return sum(scores) / len(scores)
+
+
+def _run_one_episode(
+    net: StageModule,
+    encoder: SpikeEncoder,
+    eval_data: Dataset,
+    by_class: Dict[int, List[int]],
+    embed_dim: int,
+    n_way: int,
+    n_query: int,
+    rng: random.Random,
+) -> float:
+    """Teach one memory neuron per class from one example; score queries."""
+    support_x, support_y, query_x, query_y = sample_episode(
+        eval_data, by_class, n_way, k_shot=1, n_query=n_query, rng=rng,
+    )
+    memory = OneShotAssociativeMemory(in_features=embed_dim)
+    with torch.no_grad():
+        for label in range(n_way):
+            example = support_x[support_y == label]
+            teach_from_example(memory, net, encoder.encode(example))
+        preds = memory_predictions(memory, net, encoder.encode(query_x))
+    return float((preds == query_y).float().mean())

{spikeforge-0.3.0 → spikeforge-0.3.2}/spikeforge/memory/held_out_digit_poc.py

@@ -61,12 +61,15 @@ def run_held_out_digit_poc(
     batch_size: int = 64,
     eval_batch_size: int = 256,
     train_subset: int = 1,
+    teach_examples: int = 1,
 ) -> HeldOutDigitResult:
     """Run the full protocol once and return its two accuracy numbers.
 
     ``train_subset`` (see :func:`_train_base`) never touches the
     held-out digit or either evaluation set.
     """
+    if teach_examples < 1:
+        raise ValueError("teach_examples must be positive")
     torch.manual_seed(seed)
     encoder = SpikeEncoder(coding="rate", num_steps=num_steps)
     train_data = build_dataset("mnist", train=True)
@@ -75,7 +78,9 @@ def run_held_out_digit_poc(
     net = _train_base(
         train_data, encoder, hidden, epochs, batch_size, train_subset,
     )
-    memory = _teach_held_out_digit(net, train_data, encoder)
+    memory = _teach_held_out_digit(
+        net, train_data, encoder, teach_examples,
+    )
     return _evaluate_both(
         net, memory, test_data, encoder, eval_batch_size,
     )
@@ -125,13 +130,24 @@ def _train_base(
 
 
 def _teach_held_out_digit(
-    net: StageModule, train_data: Dataset, encoder: SpikeEncoder,
+    net: StageModule,
+    train_data: Dataset,
+    encoder: SpikeEncoder,
+    teach_examples: int,
 ) -> OneShotAssociativeMemory:
-    """Build a memory module and teach it one digit-9 example."""
+    """Build memory and teach it one or more digit-9 examples."""
     memory = OneShotAssociativeMemory(in_features=hidden_size(net))
-    teach_image, _ = digit_subset(train_data, (HELD_OUT_DIGIT,))[0]
-    hidden_spikes, _ = hidden_trace(net, encoder.encode_image(teach_image))
-    memory.teach(hidden_spikes[:, 0, :])
+    held_out = digit_subset(train_data, (HELD_OUT_DIGIT,))
+    if teach_examples > len(held_out):
+        raise ValueError("teach_examples exceeds available held-out samples")
+    traces = []
+    for index in range(teach_examples):
+        teach_image, _ = held_out[index]
+        hidden_spikes, _ = hidden_trace(
+            net, encoder.encode_image(teach_image),
+        )
+        traces.append(hidden_spikes[:, 0, :])
+    memory.teach_many(torch.stack(traces))
     return memory
 
 

{spikeforge-0.3.0 → spikeforge-0.3.2}/spikeforge/memory/one_shot_associative_memory.py

@@ -76,8 +76,23 @@ class OneShotAssociativeMemory:
         the new class's index, stable for the module's lifetime and
         mapped back to a real label by the caller.
         """
+        return self.teach_many(hidden_spike_train.unsqueeze(0), gain)
+
+    def teach_many(
+        self,
+        hidden_spike_trains: torch.Tensor,
+        gain: float = DEFAULT_TEACH_GAIN,
+    ) -> int:
+        """Bind a class from several examples with one averaged write.
+
+        ``hidden_spike_trains`` is ``[N, T, H]``. Averaging the examples'
+        spike-count traces keeps this a single Hebbian write while reducing
+        sensitivity to the particular exemplar used for teaching.
+        """
+        if hidden_spike_trains.ndim != 3:
+            raise ValueError("teaching examples must have shape [N, T, H]")
         index = self._synapse.grow()
-        trace = hidden_spike_train.sum(dim=0)
+        trace = hidden_spike_trains.sum(dim=1).mean(dim=0)
         self._synapse.write(trace, index, gain)
         self._mem = None
         return index

spikeforge-0.3.2/spikeforge/memory/prototypical.py

@@ -0,0 +1,48 @@
+"""Prototypical-network embedding and loss over spiking readouts."""
+
+from typing import Tuple
+
+import torch
+
+from spikeforge.encoding.spike_encoder import SpikeEncoder
+from spikeforge.simulator.runner import run
+from spikeforge.topology.stage_module import StageModule
+
+
+def embed(
+    net: StageModule, encoder: SpikeEncoder, images: torch.Tensor,
+) -> torch.Tensor:
+    """Return ``[B, embed_dim]`` firing-rate embeddings for ``images``.
+
+    The embedding is the readout's per-neuron mean firing rate over
+    time (``Trajectory.logits``) -- the same quantity a fixed
+    classifier would softmax over, repurposed here as a spike-pattern
+    embedding rather than a class score.
+    """
+    return run(net, encoder.encode(images)).logits
+
+
+def prototypical_loss(
+    support_emb: torch.Tensor,
+    support_labels: torch.Tensor,
+    query_emb: torch.Tensor,
+    query_labels: torch.Tensor,
+    n_way: int,
+) -> Tuple[torch.Tensor, float]:
+    """Return ``(cross-entropy loss, accuracy)`` over prototype distances.
+
+    Each class's prototype is its support embeddings' mean; a query
+    is scored by negative distance to every prototype, so "close to
+    the right prototype" is trained the same way a fixed classifier
+    trains "high score for the right class".
+    """
+    prototypes = torch.stack(
+        [
+            support_emb[support_labels == c].mean(dim=0)
+            for c in range(n_way)
+        ]
+    )
+    distances = torch.cdist(query_emb, prototypes)
+    loss = torch.nn.functional.cross_entropy(-distances, query_labels)
+    accuracy = (distances.argmin(dim=1) == query_labels).float().mean()
+    return loss, float(accuracy)

spikeforge-0.3.2/spikeforge/memory/stdp_synapse.py

@@ -0,0 +1,96 @@
+"""A feedforward synapse written by pair-based spike-timing plasticity."""
+
+from typing import Optional
+
+import torch
+
+
+class STDPSynapse:
+    """Weights written step-by-step from real pre/post spike timing.
+
+    :class:`~spikeforge.memory.hebbian_synapse.HebbianSynapse` writes one
+    column in a single outer-product shot from a summed spike-count trace,
+    which is a rate code, not a timing code. This synapse instead tracks
+    exponentially-decaying eligibility traces for both sides of the
+    connection and updates every weight at every timestep from the
+    genuine order of spikes: a postsynaptic spike potentiates whichever
+    presynaptic inputs recently fired (pre-before-post, causal), and a
+    presynaptic spike depresses whichever outputs recently fired
+    (post-before-pre, anti-causal). No optimizer, no loss, no backward
+    pass touches these weights — only :meth:`step`.
+    """
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        tau_pre: float = 20.0,
+        tau_post: float = 20.0,
+        a_plus: float = 0.01,
+        a_minus: float = 0.012,
+        w_min: float = 0.0,
+        w_max: float = 1.0,
+        device: Optional[torch.device] = None,
+    ) -> None:
+        """Start from zero weights and zero eligibility traces.
+
+        ``a_minus`` slightly exceeds ``a_plus`` (the classic Song/Abbott
+        ratio) so an uncorrelated pre/post pair decays rather than drifts
+        toward saturation; ``tau_pre``/``tau_post`` set how many steps a
+        spike's influence on the other side's timing window survives.
+        """
+        self._device = device or torch.device("cpu")
+        self.tau_pre = tau_pre
+        self.tau_post = tau_post
+        self.a_plus = a_plus
+        self.a_minus = a_minus
+        self.w_min = w_min
+        self.w_max = w_max
+        self._weight = torch.zeros(
+            in_features, out_features, device=self._device
+        )
+        self._pre_trace = torch.zeros(in_features, device=self._device)
+        self._post_trace = torch.zeros(out_features, device=self._device)
+
+    @property
+    def weight(self) -> torch.Tensor:
+        """Return the current synaptic weight matrix."""
+        return self._weight
+
+    def set_weight(self, weight: torch.Tensor) -> None:
+        """Overwrite the weight matrix (e.g. to seed a starting condition)."""
+        self._weight = weight.to(self._device)
+
+    def step(
+        self, pre_spikes: torch.Tensor, post_spikes: torch.Tensor
+    ) -> None:
+        """Advance both traces by one timestep and apply the STDP update.
+
+        Trace decay happens first so a spike this step reads the *prior*
+        step's trace when computing its update — the update a spike
+        causes must depend on what came before it, not on itself.
+        """
+        pre_spikes = pre_spikes.to(self._device)
+        post_spikes = post_spikes.to(self._device)
+        with torch.no_grad():
+            decay_pre = torch.exp(torch.tensor(-1.0 / self.tau_pre))
+            decay_post = torch.exp(torch.tensor(-1.0 / self.tau_post))
+            self._pre_trace = self._pre_trace * decay_pre
+            self._post_trace = self._post_trace * decay_post
+
+            # A post spike landing on a lingering pre trace is causal
+            # (pre fired first) and potentiates; a pre spike landing on a
+            # lingering post trace is anti-causal (post fired first) and
+            # depresses. Each side reads the *other* trace, not its own.
+            potentiation = torch.outer(self._pre_trace, post_spikes)
+            depression = torch.outer(pre_spikes, self._post_trace)
+            self._weight += self.a_plus * potentiation
+            self._weight -= self.a_minus * depression
+            self._weight.clamp_(self.w_min, self.w_max)
+
+            self._pre_trace += pre_spikes
+            self._post_trace += post_spikes
+
+    def forward(self, spikes: torch.Tensor) -> torch.Tensor:
+        """Return the synaptic current ``spikes @ weight`` for one step."""
+        return spikes.to(self._device) @ self._weight