Function _drive

test/network/test_mstdp_florian.py:105–117

pre/post: [T, batch, n]; returns weight history [T, npre, npost].

(conn, pre, post, rewards)

def _drive(conn, pre, post, rewards):
    """pre/post: [T, batch, n]; returns weight history [T, npre, npost]."""
    T, batch = pre.shape[0], pre.shape[1]
    w_hist, pplus_hist, elig_hist = [], [], []
    for t in range(T):
        conn.source.s = pre[t].bool().view(batch, -1)
        conn.target.s = post[t].bool().view(batch, -1)
        conn.update(reward=rewards[t])
        w_hist.append(conn.w.detach().clone())
        ur = conn.update_rule
        pplus_hist.append(ur.p_plus.detach().clone())
        elig_hist.append(ur.eligibility.detach().clone())
    return torch.stack(w_hist), torch.stack(pplus_hist), torch.stack(elig_hist)


class TestMSTDPFlorian: