"""Spiking core: LIF reservoir + reward-modulated (three-factor) readout. Paradigm properties — no backpropagation, no gradients, no autodiff: * Leaky integrate-and-fire (LIF) neurons with refractory periods and homeostatic adaptive thresholds; * Dale's law: each reservoir neuron is permanently excitatory or inhibitory (20% inhibitory); * Readout learning is a three-factor rule: eligibility (Hebbian pre/post coincidence = reservoir trace x outcome) gated by a dopaminergic reward prediction error d = r - baseline(r); * Recurrent synapses follow a light reward-modulated STDP (trace-based), sign-constrained to respect Dale's law; * One simulation tick = one token (event-time, not wall-clock ms). All weights are constructed deterministically from a seeded mulberry32 PRNG so the TypeScript port regenerates bit-identical connectivity. """ from __future__ import annotations import numpy as np from .hashing import Mulberry32 class SpikingCore: def __init__( self, d_in: int = 1024, k_in: int = 12, n_res: int = 512, n_out: int = 4097, # 4096 vocab + 1 UNK column seed: int = 20261022, in_fanout: int = 8, rec_fanout: int = 24, inh_frac: float = 0.50, lam: float = 0.85, # membrane leak base_thr: float = 1.0, thr_min: float = 0.4, thr_max: float = 8.0, refrac: int = 2, # ticks tau_trace: float = 0.40, # reservoir spike-trace decay w_in_scale: float = 3.0, w_rec_scale: float = 0.25, w_clip: float = 4.0, eta_delta: float = 0.5, # local Rescorla-Wagner delta rate (LMS-normalized) eta_stdp: float = 0.0015, # recurrent modulated STDP rate eta_homeo: float = 0.002, # adaptive-threshold rate target_rate: float = 0.06, # target per-neuron firing probability/tick baseline_momentum: float = 0.995, # ---- hierarchical slow area G (multi-scale cortex, v2) ---- hier: bool = True, n_g: int = 256, # slow-area neurons win: int = 8, # fast ticks per slow step w_g_scale: float = 6.0, # L -> G feedforward weight scale w_recg_scale: float = 0.15, # G recurrent scale (fixed, no STDP) lam_g: float = 0.6, # slow-area membrane leak tau_trace_g: float = 0.5, # slow-area spike-trace decay target_rate_g: float = 0.10, # target firing per slow step eta_homeo_g: float = 0.01, eta_delta_g: float = 2.0, # slow-readout delta rate (8x fewer updates) eta_beta: float = 0.02, # top-down gain learning rate beta_init: float = 0.25, beta_max: float = 2.0, # ---- v3 (Muse): multi-timescale traces + direct input pathway ---- trace_taus: tuple = (0.4,), # reservoir trace decays; readout sees ALL banks direct_input: bool = False, # readout also sees a decayed trace of the # exact input-neuron pattern (ESN-standard # [state; input] readout, omitted in v1/v2) direct_tau: float = 0.5, # decay of the direct input trace direct_gain: float = 1.0, ) -> None: self.d_in, self.k_in, self.n_res, self.n_out = d_in, k_in, n_res, n_out self.seed = seed self.in_fanout, self.rec_fanout = in_fanout, rec_fanout self.inh_frac = inh_frac self.lam, self.base_thr = lam, base_thr self.thr_min, self.thr_max = thr_min, thr_max self.refrac = refrac self.tau_trace = tau_trace self.w_in_scale, self.w_rec_scale = w_in_scale, w_rec_scale self.w_clip = w_clip self.eta_delta, self.eta_stdp = eta_delta, eta_stdp self.eta_homeo, self.target_rate = eta_homeo, target_rate self.baseline_momentum = baseline_momentum self.hier = hier self.n_g, self.win = n_g, win self.w_g_scale, self.w_recg_scale = w_g_scale, w_recg_scale self.lam_g, self.tau_trace_g = lam_g, tau_trace_g self.target_rate_g, self.eta_homeo_g = target_rate_g, eta_homeo_g self.eta_delta_g = eta_delta_g self.eta_beta, self.beta0, self.beta_max = eta_beta, beta_init, beta_max # v3 structural params self.trace_taus = tuple(float(t) for t in trace_taus) self.direct_input = bool(direct_input) self.direct_tau, self.direct_gain = float(direct_tau), float(direct_gain) self.n_banks = len(self.trace_taus) self.n_state = self.n_res * self.n_banks + (self.d_in if self.direct_input else 0) # ------------------------------------------------ deterministic wiring rng = Mulberry32(seed) # Input -> reservoir: sparse fan-out per input neuron, dense storage. self.W_in = np.zeros((d_in, n_res), dtype=np.float32) for i in range(d_in): for _ in range(in_fanout): t = int(rng.next_float() * n_res) self.W_in[i, t] += rng.next_weight(w_in_scale) # Reservoir recurrent: sparse fan-out per source, Dale's law by source. self.W_rec = np.zeros((n_res, n_res), dtype=np.float32) self.inhibitory = np.zeros(n_res, dtype=np.bool_) w_rec_max = 3.0 * w_rec_scale for j in range(n_res): is_inh = rng.next_float() < inh_frac self.inhibitory[j] = is_inh sign = -1.0 if is_inh else 1.0 for _ in range(rec_fanout): t = int(rng.next_float() * n_res) self.W_rec[j, t] += sign * abs(rng.next_weight(w_rec_scale)) # Keep Dale's law under plasticity: # excitatory rows clipped to [0, w_rec_max], inhibitory to [-w_rec_max, 0] self._rec_lo = np.where(self.inhibitory, -w_rec_max, 0.0).astype(np.float32) self._rec_hi = np.where(self.inhibitory, 0.0, w_rec_max).astype(np.float32) self.W_rec = np.clip(self.W_rec, self._rec_lo[:, None], self._rec_hi[:, None]) # Readout: small seeded init, learned by the three-factor rule. rng2 = Mulberry32(seed + 1) self.W_out = np.zeros((self.n_state, n_out), dtype=np.float32) # v3: n_state rows # sparse-ish init: each column connected from a random 25% of neurons for c in range(n_out): for _ in range(n_res // 4): r = int(rng2.next_float() * n_res) self.W_out[r, c] += rng2.next_weight(0.05) # ------------------------------------------ slow area G (hierarchy) # Wired from SEPARATE PRNG streams (seed+2 / seed+3) so the fast-area # sequences above are byte-identical to v1 — the TS port depends on # the exact call order. if hier: rng3 = Mulberry32(seed + 2) # L -> G: per fast-source fanout 2. Sign follows Dale's law of the # PRESYNAPTIC fast neuron: excitatory cells excite G, inhibitory # cells inhibit it (G sees the pooled fast spike volley). self.W_inG = np.zeros((n_res, n_g), dtype=np.float32) for i in range(n_res): sign = -1.0 if self.inhibitory[i] else 1.0 for _ in range(2): t = int(rng3.next_float() * n_g) self.W_inG[i, t] += sign * abs(rng3.next_weight(w_g_scale)) # G recurrent: sparse, Dale by source, FIXED (slower areas # plasticize slower — G plasticity lives in its readout only) self.W_recG = np.zeros((n_g, n_g), dtype=np.float32) g_rec_max = 3.0 * w_recg_scale self.inhibitory_g = np.zeros(n_g, dtype=np.bool_) for j in range(n_g): is_inh = rng3.next_float() < 0.5 self.inhibitory_g[j] = is_inh sign = -1.0 if is_inh else 1.0 for _ in range(8): t = int(rng3.next_float() * n_g) self.W_recG[j, t] += sign * abs(rng3.next_weight(w_recg_scale)) self._recg_lo = np.where(self.inhibitory_g, -g_rec_max, 0.0).astype(np.float32) self._recg_hi = np.where(self.inhibitory_g, 0.0, g_rec_max).astype(np.float32) self.W_recG = np.clip(self.W_recG, self._recg_lo[:, None], self._recg_hi[:, None]) # G readout: same seeded sparse init pattern as the fast readout rng4 = Mulberry32(seed + 3) self.W_outG = np.zeros((n_g, n_out), dtype=np.float32) for c in range(n_out): for _ in range(n_g // 4): r = int(rng4.next_float() * n_g) self.W_outG[r, c] += rng4.next_weight(0.05) self.reset() # ----------------------------------------------------------------- state def reset(self) -> None: self.v = np.zeros(self.n_res, dtype=np.float32) self.thr = np.full(self.n_res, self.base_thr, dtype=np.float32) self.rhat = np.full(self.n_res, self.target_rate, dtype=np.float32) self.refr = np.zeros(self.n_res, dtype=np.int8) self.x = np.zeros(self.n_state, dtype=np.float32) # v3: augmented readout state self.x_prev = np.zeros(self.n_res, dtype=np.float32) # reservoir trace t-1 (STDP) self.x_mean = np.zeros(self.n_state, dtype=np.float32) self.x_eff = np.zeros(self.n_state, dtype=np.float32) self._banks = np.zeros((self.n_banks, self.n_res), dtype=np.float32) self._direct = np.zeros(self.d_in if self.direct_input else 0, dtype=np.float32) self.s = np.zeros(self.n_res, dtype=np.float32) # last spike vector self.baseline = 0.0 self.tick = 0 if self.hier: self.v_G = np.zeros(self.n_g, dtype=np.float32) self.thr_G = np.full(self.n_g, self.base_thr, dtype=np.float32) self.rhat_G = np.full(self.n_g, self.target_rate_g, dtype=np.float32) self.x_G = np.zeros(self.n_g, dtype=np.float32) self.x_mean_G = np.zeros(self.n_g, dtype=np.float32) self.s_G = np.zeros(self.n_g, dtype=np.float32) self.win_spikes = np.zeros(self.n_res, dtype=np.float32) self.u_slow = np.zeros(self.n_out, dtype=np.float32) self.beta = self.beta0 # ------------------------------------------------------------- snapshots def snapshot(self) -> dict: """Full dynamical state (the homeostatic scalars are LEARNED state — thresholds and the trace mean co-adapt with the readout).""" d = { "v": self.v.copy(), "thr": self.thr.copy(), "rhat": self.rhat.copy(), "refr": self.refr.copy(), "x": self.x.copy(), "x_prev": self.x_prev.copy(), "x_mean": self.x_mean.copy(), "s": self.s.copy(), "baseline": self.baseline, "banks": self._banks.copy(), "direct": self._direct.copy(), } if self.hier: d.update({ "v_G": self.v_G.copy(), "thr_G": self.thr_G.copy(), "rhat_G": self.rhat_G.copy(), "x_G": self.x_G.copy(), "x_mean_G": self.x_mean_G.copy(), "s_G": self.s_G.copy(), "win_spikes": self.win_spikes.copy(), "u_slow": self.u_slow.copy(), "beta": self.beta, }) return d def restore(self, st: dict) -> None: self.v = st["v"].copy() self.thr = st["thr"].copy() self.rhat = st["rhat"].copy() self.refr = st["refr"].copy() self.x = st["x"].copy() self.x_prev = st["x_prev"].copy() self.x_mean = st["x_mean"].copy() self.s = st["s"].copy() self.baseline = st["baseline"] if "banks" in st: self._banks = st["banks"].copy() self._direct = st["direct"].copy() else: # reconstruct banks from the composed state vector nb = self.n_banks * self.n_res self._banks = self.x[:nb].reshape(self.n_banks, self.n_res).copy() if self.direct_input: self._direct = self.x[nb:].copy() if self.hier and "v_G" in st: self.v_G = st["v_G"].copy() self.thr_G = st["thr_G"].copy() self.rhat_G = st["rhat_G"].copy() self.x_G = st["x_G"].copy() self.x_mean_G = st["x_mean_G"].copy() self.s_G = st["s_G"].copy() self.win_spikes = st["win_spikes"].copy() self.u_slow = st["u_slow"].copy() self.beta = st["beta"] def reset_dynamics(self) -> None: """Reset transient dynamics but KEEP learned homeostatic scalars (thresholds, rate estimates, trace mean, top-down gain) — used before generation.""" thr, rhat, x_mean = self.thr, self.rhat, self.x_mean if self.hier: thr_G, rhat_G, x_mean_G, beta = self.thr_G, self.rhat_G, self.x_mean_G, self.beta self.reset() self.thr, self.rhat, self.x_mean = thr, rhat, x_mean if self.hier: self.thr_G, self.rhat_G, self.x_mean_G, self.beta = thr_G, rhat_G, x_mean_G, beta # ------------------------------------------------------------- inference def step(self, input_idx) -> np.ndarray: """Advance one tick with the given active input neurons. Returns the post-synaptic trace ``x`` (the state used for readout). """ W_in, W_rec = self.W_in, self.W_rec I = W_in[input_idx].sum(axis=0, dtype=np.float32) I += self.s @ W_rec v = self.lam * self.v + I can = self.refr == 0 v[~can] = 0.0 spikes = (v >= self.thr) & can v[spikes] = 0.0 np.subtract(self.refr, 1, out=self.refr, where=self.refr > 0) self.refr[spikes] = self.refrac self.v = v self.s = spikes.astype(np.float32) # homeostatic threshold via EMA firing-rate estimate self.rhat += 0.02 * (self.s - self.rhat) self.thr += self.eta_homeo * (self.rhat - self.target_rate) np.clip(self.thr, self.thr_min, self.thr_max, out=self.thr) # traces (v3): x_prev keeps the BANK-0 (fast) reservoir trace for STDP; # the readout state x is the concatenation of all trace banks plus the # direct input trace — context older than one tick survives in the # slow banks and in the direct trace's decayed identity history. self.x_prev = self._banks[0].copy() for i, tau in enumerate(self.trace_taus): if i == 0: self._banks[i] = tau * self._banks[i] + self.s else: # slow banks are NORMALIZED traces (decay-weighted averages): # a raw tau=0.95 sum would sit at ~20x the fast bank's scale # and swamp both the readout geometry and the LMS rate. self._banks[i] = tau * self._banks[i] + (1.0 - tau) * self.s if self.direct_input: self._direct *= np.float32(self.direct_tau) np.add.at(self._direct, input_idx, np.float32(self.direct_gain)) self.x = np.concatenate([*self._banks, self._direct]) elif self.n_banks > 1: self.x = np.concatenate([*self._banks]) else: self.x = self._banks[0] # slow EMA of the trace: homeostatic "synaptic scaling" analogue used # to mean-center the readout input (kills the shared dense background # that otherwise dominates column norms) self.x_mean = 0.999 * self.x_mean + 0.001 * self.x self.x_eff = self.x - self.x_mean self.tick += 1 if self.hier: self.win_spikes += self.s return self.x def scores(self, x: np.ndarray) -> np.ndarray: """Readout scores from a (raw) trace — centers with the EMA mean.""" return (x - self.x_mean) @ self.W_out # ------------------------------------------------- slow area (hierarchy) def step_G(self) -> None: """Advance the slow area one slow step (one window of ``win`` fast ticks). Input: the pooled fast-area spike volley over the window (bottom-up, sign-preserving). Output: the cached slow readout ``u_slow``, used as a top-down prior fused into the fast prediction with gain ``beta``. Pure inference dynamics (homeostasis included, same regime as the fast area); learning happens in ``learn_G``. """ c = self.win_spikes * np.float32(1.0 / self.win) I = c @ self.W_inG I += self.s_G @ self.W_recG v = self.lam_g * self.v_G + I spikes = v >= self.thr_G v[spikes] = 0.0 self.v_G = v self.s_G = spikes.astype(np.float32) self.rhat_G += 0.05 * (self.s_G - self.rhat_G) self.thr_G += self.eta_homeo_g * (self.rhat_G - self.target_rate_g) np.clip(self.thr_G, self.thr_min, self.thr_max, out=self.thr_G) self.x_G = self.tau_trace_g * self.x_G + self.s_G self.x_mean_G = 0.999 * self.x_mean_G + 0.001 * self.x_G self.win_spikes[:] = 0.0 self.u_slow = (self.x_G - self.x_mean_G) @ self.W_outG def learn_G(self, target: int, d: float) -> int: """Delta-rule update on the SLOW readout, gated by the same dopaminergic RPE that gates the fast one (a global neuromodulator broadcasts to both areas — each readout computes its own local error from its own activity). ``target`` is the majority vocabulary column of the NEXT window. Returns 1 if the slow readout's own top prediction matches it. """ x_effG = self.x_G - self.x_mean_G u = self.u_slow V = self.n_out - 1 own_pred = int(np.argmax(u[:V])) correct = 1 if own_pred == target else 0 nx = 1.0 / (float(x_effG @ x_effG) + 1e-6) k = min(16, V) act = np.argpartition(u[:V], -k)[-k:] if target < V: cols = np.concatenate([act, np.array([target], dtype=np.int64)]) desired = np.zeros(len(cols), dtype=np.float32) desired[-1] = 1.0 else: cols = act desired = np.zeros(len(cols), dtype=np.float32) errs = u[cols] - desired gate = float(np.clip(1.0 + d, 0.25, 2.0)) eta = np.float32(self.eta_delta_g * gate * nx) W = self.W_outG W[:, cols] -= eta * errs[None, :] * x_effG[:, None] W[:, cols] = np.clip(W[:, cols], -self.w_clip, self.w_clip) return correct # ------------------------------------------------------------- learning def learn(self, pred: int, target: int, x: np.ndarray, u: np.ndarray | None = None, active_k: int = 16) -> tuple[int, float, float]: """Local delta-rule readout update (Rescorla-Wagner) gated by a dopaminergic reward-prediction error, + modulated recurrent STDP. Each readout unit computes its OWN prediction error from its own activity (no gradients, no chain rule, no backprop): units in the active top-k set and the instructive target unit move their weights to reduce (score - desired). The learning rate is gated by dopamine d = reward - baseline (RPE), the classical Rescorla-Wagner / TD account of neuromodulation. Returns (correct, reward, dopamine). """ correct = 1 if pred == target else 0 reward = float(correct) self.baseline = self.baseline_momentum * self.baseline + (1 - self.baseline_momentum) * reward d = reward - self.baseline # dopaminergic RPE W_out = self.W_out x_eff = x - self.x_mean nx = 1.0 / (float(x_eff @ x_eff) + 1e-6) # LMS normalization V = self.n_out - 1 if u is None: u = x_eff @ W_out # active set: top-k scored columns (+ target if in vocabulary) k = min(active_k, V) act = np.argpartition(u[:V], -k)[-k:] if target < V: cols = np.concatenate([act, np.array([target], dtype=np.int64)]) else: cols = act desired = np.zeros(len(cols), dtype=np.float32) if target < V: desired[-1] = 1.0 # one-hot with margin 1 errs = u[cols] - desired # dopamine gates plasticity amplitude (dips shrink, bursts amplify) gate = float(np.clip(1.0 + d, 0.25, 2.0)) eta = np.float32(self.eta_delta * gate * nx) W_out[:, cols] -= eta * errs[None, :] * x_eff[:, None] # clip in place (fancy-indexed views need explicit writeback) W_out[:, cols] = np.clip(W_out[:, cols], -self.w_clip, self.w_clip) # light reward-modulated STDP on recurrent synapses of spiked neurons if self.s.any(): spiked = np.nonzero(self.s)[0] W_rec = self.W_rec delta = np.float32(self.eta_stdp * d) # pre-trace (x_prev) paired with post spikes, gated by dopamine W_rec[:, spiked] += delta * self.x_prev[:, None] np.clip(W_rec[:, spiked], self._rec_lo[:, None], self._rec_hi[:, None], out=W_rec[:, spiked]) # top-down gain: TRUST dynamics, gated by the same RPE (a # noradrenergic gain-modulation analogue) with counterfactual credit: # fused right + fast wrong -> the slow prior SAVED the choice -> trust builds (additive) # fused wrong + fast right -> the slow prior CORRUPTED it -> trust decays (multiplicative) # Trust building slowly and decaying fast (loss aversion) lets the # system keep a genuinely informative slow area while driving a # noise- or frequency-prior-dominated one to beta ~ 0. if self.hier and target < V: pred_fast = int(np.argmax(u[:V])) if pred_fast != target and correct: # slow saved it self.beta = min(self.beta + self.eta_beta * gate, self.beta_max) elif pred_fast == target and not correct: # slow corrupted it self.beta *= float(0.98) ** gate if self.beta < 1e-4: self.beta = 0.0 else: us_t = float(self.u_slow[target]) us_p = float(self.u_slow[pred_fast]) norm = abs(us_t) + abs(us_p) + 1e-6 margin = float(np.tanh(3.0 * (us_t - us_p) / norm)) self.beta = float(np.clip( self.beta + 0.25 * self.eta_beta * gate * margin, 0.0, self.beta_max)) return correct, reward, d def learn_reservoir(self, pred: int, target: int, x: np.ndarray, u: np.ndarray | None = None) -> tuple[int, float, float]: """Training-regime dynamics WITHOUT readout plasticity: same reward, same dopaminergic RPE, same recurrent STDP — only the W_out delta update is skipped. Used at inference so the reservoir runs in its training operating regime (continual plasticity, frozen readout).""" correct = 1 if pred == target else 0 reward = float(correct) self.baseline = self.baseline_momentum * self.baseline + (1 - self.baseline_momentum) * reward d = reward - self.baseline if self.s.any(): spiked = np.nonzero(self.s)[0] W_rec = self.W_rec delta = np.float32(self.eta_stdp * d) W_rec[:, spiked] += delta * self.x_prev[:, None] W_rec[:, spiked] = np.clip(W_rec[:, spiked], self._rec_lo[:, None], self._rec_hi[:, None]) return correct, reward, d # ------------------------------------------------------- training driver def train_pass( self, neurons: np.ndarray, # (T, k_in) int32 active input neurons per tick cols: np.ndarray, # (T,) int64 target column per tick (col of token t+1) log_every: int = 5000, log_fn=None, decay_every: int = 1024, decay: float = 1e-4, record_raster_from: int | None = None, ) -> dict: """One epoch over a token stream (teacher forcing on the next token). ``neurons[t]`` drives the reservoir at tick t; the prediction made from the resulting state is scored against ``cols[t]``. """ self.reset() T = len(cols) acc_ema = None d_ema = 0.0 n_correct = 0 n_correct_fast = 0 raster: list[tuple[int, np.ndarray]] = [] V_vocab = self.n_out - 1 # UNK occupies the last column import time t0 = time.perf_counter() for t in range(T): x = self.step(neurons[t]) u = (x - self.x_mean) @ self.W_out if self.hier: # behavioural prediction = fast readout + top-down slow prior fused = u + np.float32(self.beta) * self.u_slow pred = int(np.argmax(fused[:V_vocab])) else: pred = int(np.argmax(u[:V_vocab])) # vocabulary-only prediction correct, reward, d = self.learn(pred, int(cols[t]), x, u=u) n_correct += correct if self.hier: pred_fast = int(np.argmax(u[:V_vocab])) if pred_fast == int(cols[t]): n_correct_fast += 1 if (t + 1) % self.win == 0: self.step_G() if t + 1 < T: seg = cols[t + 1 : min(t + 1 + self.win, T)] g_target = int(np.bincount(seg).argmax()) self.learn_G(g_target, d) acc_ema = reward if acc_ema is None else 0.999 * acc_ema + 0.001 * reward d_ema = 0.999 * d_ema + 0.001 * d if record_raster_from is not None and t >= record_raster_from: raster.append((t, np.nonzero(self.s)[0].astype(np.int16))) if (t + 1) % log_every == 0 and log_fn is not None: dt = time.perf_counter() - t0 log_fn({ "tick": t + 1, "total": T, "acc_ema": round(acc_ema, 4), "acc_fast_ema": (round(0.999 * (n_correct_fast / max(t + 1, 1)), 4) if self.hier else None), "beta": (round(self.beta, 3) if self.hier else None), "dopamine_ema": round(d_ema, 4), "ticks_per_sec": round((t + 1) / max(dt, 1e-9), 1), "elapsed_s": round(dt, 1), }) if (t + 1) % decay_every == 0: self.W_out *= np.float32(1.0 - decay) if self.hier: self.W_outG *= np.float32(1.0 - decay) dt = time.perf_counter() - t0 return { "ticks": T, "accuracy": n_correct / max(T, 1), "accuracy_fast_only": (n_correct_fast / max(T, 1)) if self.hier else None, "beta": float(self.beta) if self.hier else None, "acc_ema_final": acc_ema, "elapsed_s": dt, "ticks_per_sec": T / max(dt, 1e-9), "raster": raster, } # ------------------------------------------------------------ (de)hydrate def state_dict(self) -> dict: sd = { "config": { "d_in": self.d_in, "k_in": self.k_in, "n_res": self.n_res, "n_out": self.n_out, "seed": self.seed, "in_fanout": self.in_fanout, "rec_fanout": self.rec_fanout, "inh_frac": self.inh_frac, "lam": self.lam, "base_thr": self.base_thr, "refrac": self.refrac, "tau_trace": self.tau_trace, "w_in_scale": self.w_in_scale, "w_rec_scale": self.w_rec_scale, "w_clip": self.w_clip, "eta_delta": self.eta_delta, "eta_stdp": self.eta_stdp, "eta_homeo": self.eta_homeo, "target_rate": self.target_rate, "hier": self.hier, "n_g": self.n_g, "win": self.win, "w_g_scale": self.w_g_scale, "w_recg_scale": self.w_recg_scale, "lam_g": self.lam_g, "tau_trace_g": self.tau_trace_g, "target_rate_g": self.target_rate_g, "eta_homeo_g": self.eta_homeo_g, "eta_delta_g": self.eta_delta_g, "eta_beta": self.eta_beta, "beta_init": self.beta0, "beta_max": self.beta_max, "trace_taus": list(self.trace_taus), "direct_input": self.direct_input, "direct_tau": self.direct_tau, "direct_gain": self.direct_gain, }, "W_rec": self.W_rec.astype(np.float32), "W_out": self.W_out.astype(np.float32), "inhibitory": self.inhibitory, } if self.hier: sd["W_out_G"] = self.W_outG.astype(np.float32) return sd def load_state_dict(self, sd: dict) -> None: cfg = sd["config"] self.__dict__.update(cfg) # v3 derived structure (defaults reproduce v2 for old checkpoints) self.trace_taus = tuple(float(t) for t in getattr(self, "trace_taus", (self.tau_trace,))) self.direct_input = bool(getattr(self, "direct_input", False)) self.n_banks = len(self.trace_taus) self.n_state = int(sd["W_out"].shape[0]) self.W_rec = np.asarray(sd["W_rec"], dtype=np.float32) self.W_out = np.asarray(sd["W_out"], dtype=np.float32) self.inhibitory = np.asarray(sd["inhibitory"], dtype=np.bool_) w_rec_max = 3.0 * self.w_rec_scale self._rec_lo = np.where(self.inhibitory, -w_rec_max, 0.0).astype(np.float32) self._rec_hi = np.where(self.inhibitory, 0.0, w_rec_max).astype(np.float32) if "W_out_G" in sd and self.hier: self.W_outG = np.asarray(sd["W_out_G"], dtype=np.float32) else: # v1 checkpoint without the slow area self.hier = False self.reset()