PC-ALM

mujou.pc

PC-ALM (augmented Lagrangian predictive coding) as nodes of mujou.

Each node holds a part of the model and learns from its neighbors' messages only: Input takes the samples the host writes, each Layer wraps a module as one constraint h = f(h_in), and Output holds the last module and the loss. Inside a node, the module is differentiated by autograd; between nodes, activities (Activity, the predictions of the forward pass) flow to the next node and gradients of the energy to the previous one. The target of each sample travels with its input in the Activity, so the output node pairs every prediction with its own target. Chain connects the nodes for a sequence of modules, and show shows it a sample. Each node gives the spec of what it sends as act_spec, which the next node takes, to build other shapes by hand.

The names that go with a role are in a submodule, which this package imports:

• pc.optim: SGD and AdamW set how the parameters of a node move in time (Optimizer); each node binds the setting to its own parameters (BoundOptimizer).
• pc.relax: Euler sets how the deviation r = h - mu of a node moves (Relaxer); each node binds the setting to its own r (BoundRelaxer).
• pc.loss: MSELoss and CrossEntropyLoss give the loss of each sample.
• pc.msg: Sample, Activity, Prediction and activity_spec, what the nodes send.
• pc.clock: LearnClock and Tick.
• pc.functional: the parts the nodes are made of, for writing nodes of your own.

The samples are independent and the shapes start with the batch, batch: () (the default) for one sample at a time, B for B samples, or a shape of several dimensions.

Import it as a submodule::

import mujou.pc as pc

B = 64
net = pc.Chain(
    nn.Sequential(nn.Linear(784, 256), nn.ReLU()),
    nn.Linear(256, 10),
    x=TensorSpec((B, 784)),
    y=TensorSpec((B,), torch.int64),
    batch=B,
    loss=pc.loss.CrossEntropyLoss(B),
    tau=pc.TimeConstants(h=4e-3, lam=2e-3, r=4e-3),
    optimizer=pc.optim.SGD(tau=1.28),
    learn_interval=0.15,
)

Chain

Chain(
    *modules: Module,
    x: TensorSpec,
    y: TensorSpec,
    loss: Callable[[Tensor, Tensor], Tensor],
    tau: TimeConstants,
    optimizer: Optimizer,
    relaxer: Relaxer = Euler(),
    learn_interval: float | None,
    batch: Size = (),
    max_dt: float | None = None,
    period: float | None = None,
    device: Device = "cpu",
)

Bases: Network, Shown

Input -> Layer -> ... -> Output: a network that learns modules in sequence by PC-ALM, as nn.Sequential(*modules) would compute them.

modules[:-1] each become a Layer and modules[-1] the Output, after an Input of the samples x and targets y. Each node takes the act_spec of the one before and the nodes are connected in order. The children are input, layer0, ..., output (modules[i] is in layer{i}, also layers[i]), which are also their paths in Run.state_dict(). Every node runs with the same period, tau.period (tau.h / 2 with tau.r = tau.h) unless given; for a period of its own for each node, build the nodes one by one (show shows them samples too).

The host writes a Sample to prev and reads a Prediction from next, or does both with show.

Example::

B = 64
net = pc.Chain(
    nn.Sequential(nn.Linear(784, 256), nn.ReLU()),
    nn.Linear(256, 10),
    x=TensorSpec((B, 784)),
    y=TensorSpec((B,), torch.int64),
    batch=B,
    loss=pc.loss.CrossEntropyLoss(B),
    tau=pc.TimeConstants(h=4e-3, lam=2e-3, r=4e-3),
    optimizer=pc.optim.SGD(tau=1.28),
    learn_interval=0.15,
)
with mujou.run(net) as run:
    prediction = pc.show(run, net, sample, seconds=0.2)

PARAMETER	DESCRIPTION
modules	The functions of the nodes, in order, on device. TYPE: Module DEFAULT: ()
x	The spec of the inputs, (*batch, *in). TYPE: TensorSpec
y	The spec of the targets, (*batch, *target). TYPE: TensorSpec
loss	The loss of each sample, as in Output. TYPE: Callable[[Tensor, Tensor], Tensor]
tau	The time constants of every node. TYPE: TimeConstants
optimizer	How the parameters of every node move, as in Layer: each node binds it to its own parameters. TYPE: Optimizer
relaxer	How r of every Layer moves, as in Layer: each binds it to its own r. The Output has no r. TYPE: Relaxer DEFAULT: Euler()
learn_interval	As in Layer. Longer than the round trip of the credit, 2 * len(modules) * period, is recommended. TYPE: float | None
batch	The batch, as in Input: () for one sample, B for B samples. TYPE: Size DEFAULT: ()
max_dt	As in Layer. TYPE: float | None DEFAULT: None
period	The period of every node, tau.period if None: with the defaults of TimeConstants, the credit moves one node down in each period. A period of your own must not be longer than max_dt. TYPE: float | None DEFAULT: None
device	Where the nodes compute. TYPE: Device DEFAULT: 'cpu'

RAISES	DESCRIPTION
ValueError	As Input, Layer and Output raise.

input instance-attribute

input = Input(x=x, y=y, period=period, batch=batch, device=device)

output instance-attribute

output = Output(
    modules[-1],
    act_spec=last.act_spec,
    loss=loss,
    tau=tau,
    optimizer=optimizer,
    learn_interval=learn_interval,
    period=period,
    max_dt=max_dt,
    device=device,
)

prev instance-attribute

prev = self.input.prev

next instance-attribute

next = self.output.next

layers property

layers: tuple[Layer, ...]

The Layer of each module but the last, in order: layers[i] is the child layer{i}.

Shown

Bases: ABC

What show shows a sample: a network whose prev takes a Sample and whose next sends a Prediction.

Chain is one. A network built by hand is one by inheriting this class as well as Network, and assigning the prev of its Input and the next of its Output to prev and next in __init__::

class Tower(mujou.Network, pc.Shown):
    def __init__(self) -> None:
        super().__init__()
        ...
        self.prev, self.next = self.input.prev, self.output.next

prev instance-attribute

prev: Port[Never, Sample]

next instance-attribute

next: Port[Prediction, Never]

Input

Input(
    *,
    x: TensorSpec,
    y: TensorSpec,
    period: float,
    batch: Size = (),
    device: Device = "cpu",
)

Bases: Node

A node that takes the samples the host writes and sends them on as the Activity of the next node, with their targets.

The host writes a Sample (x, y and mask) to prev with Run.write. In each on_update the node sends from next an Activity whose h is x, with y, mask and the time (ns) at which the host wrote them: the nodes on the path pass them on with the activity computed from x, so that the output node pairs each prediction with its own target. A new time tells them the input changed. x is clamped: what the next node sends back on next is not read.

Before the host first writes, the node sends zeros with mask False and time 0. It holds no state of its own: what it sends is the last write. act_spec is the spec of what the node sends, for the next node.

Example::

B = 64
inp = pc.Input(
    x=TensorSpec((B, 784)),
    y=TensorSpec((B,), torch.int64),
    period=1e-3,
    batch=B,
)
...
mask = torch.ones(B, dtype=torch.bool)
with mujou.run(net) as run:
    run.write({inp.prev: {"x": x, "y": y, "mask": mask}})

PARAMETER	DESCRIPTION
x	The spec of the inputs x, (*batch, *in). TYPE: TensorSpec
y	The spec of the targets y, (*batch, *target). TYPE: TensorSpec
period	The period of the node in seconds, usually that of the next node. TYPE: float
batch	The batch that x and y start with, as in activity_spec: () for one sample, B for B samples. TYPE: Size DEFAULT: ()
device	Where the node computes. TYPE: Device DEFAULT: 'cpu'

RAISES	DESCRIPTION
ValueError	x or y does not start with batch.

act_spec instance-attribute

act_spec = activity_spec(x, y, batch)

prev instance-attribute

prev = Port(recv=DictSpec(Sample, x=x, y=y, mask=mask, reset=mask))

next instance-attribute

next = MultiPort(send=self.act_spec, recv=x)

on_update

on_update(dt: float) -> None

Layer

Layer(
    module: Module,
    *,
    act_spec: DictSpec[Activity],
    tau: TimeConstants,
    optimizer: Optimizer,
    relaxer: Relaxer = Euler(),
    learn_interval: float | None,
    period: float | None = None,
    max_dt: float | None = None,
    device: Device = "cpu",
)

Bases: Node

A node that wraps module as one constraint h = f(h_in; theta) of PC-ALM.

It predicts mu = module(h_in) from the activity h_in of the previous node and sends it on: the activity of every node is the forward pass. It holds the deviation \(r = h - \mu\) of its activity and the dual lam of the constraint, which stay in the node and make the credit it sends back. With \(\delta = r + \lambda\) and the sum \(g\) of what the next nodes send back:

\[ \tau_h \dot r = -(r + \lambda + g) - \frac{\tau_h}{\tau_r}\,r,\qquad \tau_\lambda \dot\lambda = r, \]

and theta moves by optimizer (such as pc.optim.SGD, \(\tau_\theta \dot\theta = -g_\theta / B\)) with the mean gradient \(g_\theta / B\), where \(B\) is the number of samples in the batch (1 without one), \(g_\theta = -J_\theta^\top \delta\) and the node sends \(g_{\mathrm{in}} = -J_h^\top \delta\) to the previous node (both from one vector-Jacobian product of module, functional.vjp). At a fixed point, \(r = 0\) and lam is the credit \(-\partial\ell/\partial h\) of backpropagation. Inside the node, module is differentiated by autograd; between nodes, only these messages flow.

Sending mu rather than mu + r is the fixed prediction assumption of predictive coding: it keeps the fixed point, and the credit flows down the path only. Were r sent up, it would move the activities of the nodes above and come back as credit through the product of their Jacobians, a loop that diverges in chains whose forward pass amplifies, however short the steps. With the defaults of TimeConstants (period = tau.h / 2, tau.lam = tau.h / 2, tau.r = tau.h), delta follows -g in one update, so the credit moves one node down in each period.

Whatever moves mu (a new sample, an update of theta here or below), r and lam stay and only relax, unless the sample starts anew: in the update that first sees an input, the node sets r and lam of the samples whose reset is set to zero, so that no credit of an unrelated sample is left (Sample). Without a target, r and lam decay to 0.

Each on_update(dt):

1. reads the activity and ticks its clock, which caps dt at max_dt; with a new input, sets r and lam of the samples it resets to zero;
2. computes mu, g_in and, when it learns, g_theta with the delta before this update, zero for the samples whose mask is False, and sends g_in;
3. moves r and then lam by dt with the g the next nodes sent (functional.relax_);
4. sends h = mu, and the y, mask, reset and time that came with h_in;
5. when it learns, moves theta by the time since its last update (optimizer.step_).

Its clock (a LearnClock) tells when it learns: in its first update once learn_interval seconds have passed since the host wrote the input (its time, on the run's time), and then every learn_interval seconds while the input stays, but never in the update that first sees an input (its delta is still that of the input before); with learn_interval=None, in every update; with math.inf, never. The nodes read the one timer of the run, so the nodes on a path learn at about the same time, however deep they are, and each from the relaxed delta of the same input. A learn_interval longer than the round trip of the credit down the path and back (2 * hops * period, with hops the number of connections from the Input to the Output) is recommended, so that the nodes learn from the relaxed delta. The node does not learn when no sample of the input has a target (every mask is False): that time is dropped from the next update, which would otherwise learn from the delta left over from an earlier input. Nor does it keep the time of an input that changed before the node learned from it. Until the first sample reaches the node, the activity is zeros (mask all False, time 0), so it does not learn then.

dt is capped because the step is explicit: after a long stall (a loaded machine), a long dt / tau.h would overshoot and make r and lam diverge. The time beyond max_dt is dropped, from the relaxation and from the weight of the next update of theta alike. The default tau.max_dt is the step that settles delta (tau.h / 2 with tau.r = tau.h).

relaxer sets how r moves along the force on it (functional.relax_): the plain step of pc.relax.Euler, or a relaxer of your own with the same fixed point. The node calls the reset_ of the bound relaxer with each new input, with the samples it resets (whose state the relaxer forgets).

module must treat the samples of the batch independently (no BatchNorm in training mode): the gradient of each sample must be its own. It must also be deterministic (no Dropout): mu is what the node sends. Its parameters that require gradients learn; set requires_grad to False to freeze one. It is called once on zeros when the node is made, to find the shape and dtype of h.

The state (Run.state_dict()["nodes"][path]["state"]) holds "generator" (the state of the node's generator), "module" (its state_dict), "r" and "lam" ((*batch, *out)), when to learn ("clock", the state of the LearnClock), the state of optimizer ("optimizer", such as the moments of pc.optim.AdamW) and that of relaxer ("relaxer", empty for pc.relax.Euler). The h it sent last is in the state of its port next. To watch the node while it runs, read them there: Run.state_dict() pauses the run, whereas reading module, r or lam from the host's thread races with the node's updates in place. act_spec is the spec of what the node sends, for the next node.

Example::

B = 64
tau = pc.TimeConstants(h=4e-3, lam=2e-3, r=4e-3)
inp = pc.Input(
    x=TensorSpec((B, 784)),
    y=TensorSpec((B,), torch.int64),
    period=tau.period,
    batch=B,
)
hidden = pc.Layer(
    nn.Sequential(nn.Linear(784, 256), nn.ReLU()),
    act_spec=inp.act_spec,
    tau=tau,
    optimizer=pc.optim.SGD(tau=1.28),
    learn_interval=24e-3,
)
inp.next.connect(hidden.prev)
hidden.next.connect(out.prev)  # out: a pc.Output, or another pc.Layer

PARAMETER	DESCRIPTION
module	The function f, (*batch, *in) -> (*batch, *out), on device. TYPE: Module
act_spec	The spec of the activity h_in it receives: the act_spec of the node before, which also gives the batch. TYPE: DictSpec[Activity]
tau	The time constants of r and lam. TYPE: TimeConstants
optimizer	How theta moves (pc.optim.SGD, pc.optim.AdamW or an Optimizer of your own). The node binds it to its parameters that require gradients and holds the bound optimizer as optimizer. TYPE: Optimizer
relaxer	How r moves along the force on it (pc.relax.Euler, the default, or a Relaxer of your own). The node binds it to its r and holds the bound relaxer as relaxer. TYPE: Relaxer DEFAULT: Euler()
learn_interval	The time in seconds from the write of an input to the update of theta, and between the updates while the input stays; longer than 2 * hops * period is recommended. None updates in every on_update, for a stream whose input changes at every frame (good values for it have not been measured yet); math.inf never updates. TYPE: float | None
period	The period of the node in seconds; None is tau.period (tau.h / 2 with tau.r = tau.h). It must not be longer than max_dt. TYPE: float | None DEFAULT: None
max_dt	The longest dt an update counts, in seconds; None is tau.max_dt. TYPE: float | None DEFAULT: None
device	Where the node computes. TYPE: Device DEFAULT: 'cpu'

RAISES	DESCRIPTION
ValueError	The output of module does not start with the batch of act_spec, or LearnClock raises: learn_interval is neither None nor > 0, max_dt is not > 0, or period is longer than max_dt. module raises as PyTorch does when it does not take the h of act_spec on device.

clock instance-attribute

clock = LearnClock.of(
    tau, period=self.period, learn_interval=learn_interval, max_dt=max_dt
)

module instance-attribute

module = module

optimizer instance-attribute

optimizer = optimizer.bind(self._params)

r instance-attribute

r = torch.zeros_like(mu)

relaxer instance-attribute

relaxer = relaxer.bind(self.r)

lam instance-attribute

lam = torch.zeros_like(mu)

act_spec instance-attribute

act_spec = activity_spec(out, act_spec.field('y', TensorSpec), batch)

prev instance-attribute

prev = Port(send=h_in, recv=act_spec)

next instance-attribute

next = MultiPort(send=self.act_spec, recv=out)

on_update

on_update(dt: float) -> None

Output

Output(
    module: Module,
    *,
    act_spec: DictSpec[Activity],
    loss: Callable[[Tensor, Tensor], Tensor],
    tau: TimeConstants,
    optimizer: Optimizer,
    learn_interval: float | None,
    period: float | None = None,
    max_dt: float | None = None,
    device: Device = "cpu",
)

Bases: Node

A node that predicts y_hat = module(h_in) and whose energy is the loss of the samples that have a target.

It pairs h_in with the target y and mask of the same Activity, so the prediction always meets the target of the very input it was computed from, however fast the host writes. With

\[ \ell = \sum_b \mathrm{mask}_b\, \ell_b(\hat y_b, y_b), \]

theta moves by optimizer with the mean gradient \(\frac{1}{B}\frac{\partial\ell}{\partial\theta}\) (pc.optim.SGD makes it \(\tau_\theta \dot\theta = -\frac{1}{B}\frac{\partial\ell}{\partial\theta}\)), where \(B\) is the number of samples in the batch (1 without one), and it sends \(g_{\mathrm{in}} = \partial\ell/\partial h_{\mathrm{in}}\) to the previous node (functional.loss_vjp): zero for a sample without a target while its y_hat is finite. Each on_update(dt) caps dt at max_dt as Layer does, reads the Activity, sends g_in, sends a Prediction (y_hat with the y, mask and time it was paired with) for the host to read, and updates theta when it learns.

It learns as Layer does, by its clock: once learn_interval seconds have passed since the host wrote the input, but not in the update that first sees it, or in every update with None. When no sample has a target, it does not learn and drops the time since the last update; so too before the first sample arrives, while the activity is zeros (mask all False), and when the input changes before it has learned from it. B is the number of samples in the batch, not the number of samples with a target: a sample without one adds zero.

loss returns the loss of each sample, of the shape of the batch (that of mask), such as MSELoss(batch) or CrossEntropyLoss(batch), and the node sums them. A loss that averages over the batch (such as nn.MSELoss()) would scale the gradient that flows back by 1 / B, so the node refuses it when it is made: it calls module and loss once on zeros then (functional.loss_vjp), which also finds the shape of y_hat. module must treat the samples independently, as in Layer.

The state holds "generator", "module", when to learn ("clock") and the state of optimizer ("optimizer"). Read it with Run.state_dict(), as for Layer.

Example::

out = pc.Output(
    nn.Linear(256, 10),
    act_spec=hidden.act_spec,  # hidden: a pc.Layer
    loss=pc.loss.CrossEntropyLoss(B),
    tau=tau,
    optimizer=pc.optim.SGD(tau=1.28),
    learn_interval=24e-3,
)
hidden.next.connect(out.prev)
# The host: prediction = run.read(out.next); prediction["y_hat"], ...

PARAMETER	DESCRIPTION
module	The function f, (*batch, *in) -> (*batch, *out), on device. TYPE: Module
act_spec	The spec of the activity h_in it receives: the act_spec of the node before, which also gives the batch. TYPE: DictSpec[Activity]
loss	The loss of each sample, (y_hat, y) -> batch. TYPE: Callable[[Tensor, Tensor], Tensor]
tau	The time constants; the defaults of period and max_dt are taken from them. TYPE: TimeConstants
optimizer	How theta moves, as in Layer. TYPE: Optimizer
learn_interval	As in Layer. TYPE: float | None
period	The period of the node in seconds; None is tau.period (tau.h / 2 with tau.r = tau.h). It must not be longer than max_dt. TYPE: float | None DEFAULT: None
max_dt	As in Layer: the longest dt an update counts, here towards the weight of the update of theta. The default tau.max_dt is that of Layer, so that every node of a chain counts the same time. TYPE: float | None DEFAULT: None
device	Where the node computes. TYPE: Device DEFAULT: 'cpu'

RAISES	DESCRIPTION
ValueError	LearnClock raises, as in Layer, or loss does not return the shape of the batch for the output of module and the target.

clock instance-attribute

clock = LearnClock.of(
    tau, period=self.period, learn_interval=learn_interval, max_dt=max_dt
)

module instance-attribute

module = module

optimizer instance-attribute

optimizer = optimizer.bind(self._params)

prev instance-attribute

prev = Port(send=h_in, recv=act_spec)

next instance-attribute

next = Port(
    send=DictSpec(
        Prediction,
        y_hat=TensorSpec(y_hat.shape, y_hat.dtype),
        y=y,
        mask=mask,
        time=act_spec.field("time", TimestampSpec),
    )
)

on_update

on_update(dt: float) -> None

TimeConstants dataclass

TimeConstants(*, h: float, lam: float, r: float)

The time constants of the continuous-time PC-ALM, in seconds.

A node sends its prediction mu and holds the deviation \(r = h - \mu\) of its activity from it and the dual lam, which it relaxes with

\[ \tau_h \dot r = -(r + \lambda + g) - \frac{\tau_h}{\tau_r}\,r,\qquad \tau_\lambda \dot\lambda = r, \]

and a step of dt seconds moves them by dt / tau of the right-hand side (the discrete step sizes are eta_h = dt / h and alpha = dt / lam). The parameters theta move by the optimizer of the node (pc.optim.SGD or pc.optim.AdamW), which holds their time constant.

lam=math.inf makes the dual stay at zero (alpha = 0); with r=math.inf too, the nodes are those of predictive coding with the fixed prediction assumption, whose error r settles at the credit of backpropagation (the decay would shrink it by 1 + h / r at each node). h must be finite.

r is the time constant of the decay of the deviation. The fixed point stays (lam integrates r, so r = 0 there, and lam is the credit of backpropagation). The nodes send mu, so the credit flows down the path only, and the steps are stable in any chain once they are in each node. The defaults are chosen together: with r = h, lam = h / 2 and a period of h / 2, one update moves delta = r + lam to -g exactly, so the credit moves one node in each period and does not overshoot (a larger step overshoots, and the overshoot grows at every node). An analog integrator that holds r and leaks is this decay.

The fields are keyword-only and have no defaults: all are floats, and a swap by position could not be caught by the type checker.

Example::

tau = TimeConstants(h=4e-3, lam=2e-3, r=4e-3)  # a period of 2e-3
pc_only = TimeConstants(h=4e-3, lam=math.inf, r=math.inf)

RAISES	DESCRIPTION
ValueError	h is not finite and > 0, or lam or r is not > 0 (NaN is not).

h instance-attribute

h: float

lam instance-attribute

lam: float

r instance-attribute

r: float

period property

period: float

The default period of a node, h / (1 + h / r): the step that moves r all the way along its own pull (h / 2 with r = h, h without the decay). With lam = h * r / (h + r) (h / 2 with r = h), it moves delta to -g in one update.

max_dt property

max_dt: float

The default longest dt an update counts, min(h / (1 + h / r), 2 * lam) (the default period unless the dual is faster): the explicit Euler step of r then lam is then stable (specs/pc.md, "dt の打ち切り"). Every node of a chain caps at the same time, so that they weigh their updates of the parameters alike.

show

show(run: Run, chain: Shown, sample: Sample, *, seconds: float) -> Prediction

Return the prediction chain makes for sample before it learns from it (a prequential prediction), after showing it seconds long.

chain is a Shown: a Chain, or a network of the nodes of mujou.pc built by hand that inherits Shown. Writes sample to chain.prev, waits for the first prediction the host reads for that write, and returns it once seconds (of run.timer) have passed since the write, while the nodes learn from the sample (Sample.mask True) or only relax towards it (every mask False, as to evaluate, where a seconds that lets the input reach the output is enough). A node learns from a sample from its second update with it on (learn_interval after the write, when set), and the first prediction of the sample is made of the first activities with it, so the prediction comes before learning unless the host reads it later than that.

Nothing moves the timer of a ManualRun, so there it does not return.

Example::

mask = torch.ones(B, dtype=torch.bool)
sample = pc.msg.Sample(x=x, y=y, mask=mask, reset=mask)  # shuffled
p = pc.show(run, net, sample, seconds=0.2)
correct = int((p["y_hat"].argmax(1) == y).sum())

RAISES	DESCRIPTION
ValueError	A tensor of sample has another shape or dtype than the spec of chain.prev.
TypeError	sample lacks a key, or has one too many.
RuntimeError	The run shut down normally before the prediction.
NetworkError	The run stopped because a node raised.

mujou.pc.optim

The optimizers that move the parameters of a node in time: SGD and AdamW.

An optimizer here is a setting (Optimizer): a frozen value that holds no state and may be shared by every node. A node binds it to its own parameters with bind(params), which returns the BoundOptimizer that holds the state (the moments of AdamW) and moves the parameters by step_(grads, elapsed). The time constants are in seconds, so the flow does not depend on how often the node learns: elapsed is the time the step weighs.

They do not use torch.optim: its optimizers count steps rather than time, and making one brings the GIL back into a free-threaded Python.

Optimizer

Bases: ABC

The setting of how the parameters of a node move in time.

It holds values only, no state, so the nodes may share it. bind(params) makes the BoundOptimizer of the parameters of one node, which holds the state. Give it only the parameters that learn (pc.Layer and pc.Output bind those that require gradients).

An optimizer of your own inherits this class, and its bound optimizer BoundOptimizer. For example, SGD with a learning rate per step rather than per second::

@dataclass(frozen=True)
class StepSGD(pc.optim.Optimizer):
    lr: float

    def bind(self, params):
        return _BoundStepSGD(list(params), self.lr)

class _BoundStepSGD(pc.optim.BoundOptimizer):
    def __init__(self, params, lr):
        self._params, self._lr = params, lr

    def step_(self, grads, elapsed):
        with torch.no_grad():
            for p, g in zip(self._params, grads, strict=True):
                p.sub_(g, alpha=self._lr)

    def state_dict(self):
        return {}

    def load_state_dict(self, state):
        pass

bind abstractmethod

bind(params: Iterable[Tensor]) -> BoundOptimizer

Return the optimizer bound to params, with a fresh state.

BoundOptimizer

Bases: Stateful

An optimizer bound to the parameters of one node, which holds its state.

It holds the parameters by reference. The node that bound it owns it and calls it from its own hooks only. It is a Stateful, so an attribute of a node that holds it is in the node's state by itself: state_dict returns the state with tensors that share memory with the optimizer's, and load_state_dict copies the tensors of the state into those (it keeps none of the state it is given).

step_ abstractmethod

step_(grads: Sequence[Tensor], elapsed: float) -> None

Move the parameters in place by elapsed seconds of the flow, without recording gradients.

PARAMETER	DESCRIPTION
grads	The gradients of the parameters, in the order they were bound and of their shapes, dtypes and devices, averaged over the batch (the node applies the 1 / B; the optimizer does not know the batch). TYPE: Sequence[Tensor]
elapsed	The time the step weighs, in seconds (Tick.elapsed: the time since the parameters last moved, capped, without the time the clock dropped). TYPE: float

RAISES	DESCRIPTION
ValueError	grads does not have one gradient per parameter.

state_dict abstractmethod

state_dict() -> Any

Return the state. Its tensors may share memory with the object's: the run copies the state it takes.

load_state_dict abstractmethod

load_state_dict(state: Any) -> None

Restore the state that state_dict returned, in place.

SGD dataclass

SGD(*, tau: float)

Bases: Optimizer

Gradient descent in time, \(\tau_\theta \dot\theta = -g\).

A step of elapsed seconds is one step of SGD with the learning rate elapsed / tau: p -= (elapsed / tau) * g. With \(\Delta t\) seconds between two updates of the parameters (the time each input is shown, for a node that learns once from each), the step is of the learning rate Δt / tau. The bound optimizer has no state (an empty state_dict).

Example::

hidden = pc.Layer(module, act_spec=..., tau=tau,
                  optimizer=pc.optim.SGD(tau=0.3), learn_interval=24e-3)

PARAMETER	DESCRIPTION
tau	The time constant \(\tau_\theta\) of the parameters, in seconds. TYPE: float

RAISES	DESCRIPTION
ValueError	tau is not finite and > 0. Never to learn, give the node learn_interval=math.inf.

tau instance-attribute

tau: float

bind

bind(params: Iterable[Tensor]) -> BoundOptimizer

Return the optimizer bound to params.

AdamW dataclass

AdamW(
    *,
    tau: float,
    tau_m: float,
    tau_v: float,
    eps: float = 1e-08,
    weight_decay: float = 0.01,
)

Bases: Optimizer

AdamW in time: moments that decay with time constants, and decoupled weight decay.

With the gradient \(g\) (averaged over the batch):

\[ \tau_m \dot m = g - m,\qquad \tau_v \dot v = g^2 - v,\qquad \tau_\theta \dot\theta = -\left(\frac{\hat m}{\sqrt{\hat v} + \varepsilon} + w\,\theta\right), \]

where \(\hat m = m / c_m\) and \(\hat v = v / c_v\) correct the bias of the moments that start at zero, with the masses \(\tau_m \dot c_m = 1 - c_m\) and \(\tau_v \dot c_v = 1 - c_v\) from \(c = 0\): \(c\) is the share of the moment made of gradients so far (the rest is the zero it started from), \(1 - e^{-t/\tau}\) for gradients seen over the time \(t\).

A step of elapsed \(= \Delta t\) seconds moves the moments and masses exactly as the flow would with \(g\) held over \(\Delta t\) (\(m \leftarrow m + a_m (g - m)\) and \(c_m \leftarrow c_m + a_m (1 - c_m)\) with \(a_m = 1 - e^{-\Delta t/\tau_m}\), the same for \(v\)), so that two steps of \(\Delta t\) and one of \(2\Delta t\) give the same moments for the same \(g\). \(\theta\) then steps with \(\eta = \Delta t / \tau_\theta\), in the order of torch.optim.AdamW: the decay first, on the \(\theta\) before the step, and the Adam term with the moments and masses just moved:

\[ \theta \leftarrow \theta (1 - \eta w) - \eta\,\frac{\hat m}{\sqrt{\hat v} + \varepsilon}. \]

A step of elapsed == 0 moves nothing.

With a constant \(\Delta t\), \(1 - a = e^{-\Delta t/\tau}\) is the \(\beta\) of torch.optim.AdamW and \(c = 1 - \beta^t\) its bias correction after \(t\) steps: the steps are those of torch.optim.AdamW(lr=Δt / tau, betas=(exp(-Δt / tau_m), exp(-Δt / tau_v)), eps=eps, weight_decay=weight_decay), up to rounding. From the \(\beta\) of a node that learns every \(\Delta t\), \(\tau_m = -\Delta t / \ln\beta_1\): learning every 32 ms, \(\beta_1 = 0.9\) is about 0.30 s and \(\beta_2 = 0.999\) about 32 s.

weight_decay \(w\) has no unit, and decays \(\theta\) with the time constant \(\tau_\theta / w\), as weight_decay does with lr in torch.optim.AdamW.

The state (state_dict()) holds the moments "m" and "v" (a list of tensors each, of the dtype and device of the parameters) and the masses "c_m" and "c_v" (Python floats in [0, 1]). There is no count of steps.

A step moves each element by about \(\eta\) whatever the size of its gradient, so tau is Δt / lr for the lr of Adam, with \(\Delta t\) the time between two updates of the parameters.

Example (lr = 1e-3, learning once from each input, shown for 32 ms)::

hidden = pc.Layer(module, act_spec=..., tau=tau,
                  optimizer=pc.optim.AdamW(tau=32e-3 / 1e-3, tau_m=0.3,
                                     tau_v=30.0),
                  learn_interval=24e-3)

PARAMETER	DESCRIPTION
tau	The time constant \(\tau_\theta\) of the parameters, in seconds. TYPE: float
tau_m	The time constant \(\tau_m\) of the first moment, in seconds. TYPE: float
tau_v	The time constant \(\tau_v\) of the second moment, in seconds. TYPE: float
eps	The \(\varepsilon\) added to \(\sqrt{\hat v}\). It is > 0, so that an element whose gradient has been 0 does not become 0 / 0. TYPE: float DEFAULT: 1e-08
weight_decay	The decoupled weight decay \(w\). TYPE: float DEFAULT: 0.01

RAISES	DESCRIPTION
ValueError	tau, tau_m, tau_v or eps is not finite and 0, or weight_decay is not finite and >= 0.

tau instance-attribute

tau: float

tau_m instance-attribute

tau_m: float

tau_v instance-attribute

tau_v: float

eps class-attribute instance-attribute

eps: float = 1e-08

weight_decay class-attribute instance-attribute

weight_decay: float = 0.01

bind

bind(params: Iterable[Tensor]) -> BoundOptimizer

Return the optimizer bound to params, with the moments and the masses at zero.

mujou.pc.relax

The relaxers that move the deviation r = h - mu of a node in time: Euler, or one of your own.

A relaxer is a setting (Relaxer), as an optimizer is: a frozen value that holds no state and may be shared by every node. A node binds it to its own deviation with bind(r), which returns the BoundRelaxer that holds the state (if the relaxer has one) and moves r by step_(r, force, dt, tau_h). It is the counterpart of pc.optim.SGD and pc.optim.AdamW for r: those move the parameters along their gradient, a relaxer moves the deviation along the force on it, \(F = (1 + \tau_h / \tau_r)\,r + \lambda + g\) (functional.relax_). The dual lam is not a relaxer's business: it follows the moved r.

A relaxer changes how r gets to its fixed point, not where it is: at the fixed point the force is zero and the state of the relaxer is at rest, and r stays there.

Relaxer

Bases: ABC

The setting of how the deviation r = h - mu of a node moves in time.

It holds values only, no state, so the nodes may share it. bind(r) makes the BoundRelaxer of the deviation of one node, which holds the state.

A relaxer of your own inherits this class, and its bound relaxer BoundRelaxer. For example, a step that is capped in size::

@dataclass(frozen=True)
class Capped(pc.relax.Relaxer):
    limit: float

    def bind(self, r):
        return _BoundCapped(self.limit)

class _BoundCapped(pc.relax.BoundRelaxer):
    def __init__(self, limit):
        self._limit = limit

    def step_(self, r, force, dt, tau_h):
        step = (dt / tau_h) * force
        r.sub_(step.clamp(-self._limit, self._limit))

    def reset_(self, samples):
        pass

    def state_dict(self):
        return {}

    def load_state_dict(self, state):
        pass

Its fixed point must be a zero of the force (with the state of the relaxer at rest), or the node no longer solves PC-ALM: lam is the credit of backpropagation only where r stops at a zero of the force.

bind abstractmethod

bind(r: Tensor) -> BoundRelaxer

Return the relaxer bound to a deviation of the shape, dtype and device of r, with a fresh state.

BoundRelaxer

Bases: Stateful

A relaxer bound to the deviation r of one node, which holds its state.

The node that bound it owns it and calls it from its own hooks only. It is a Stateful, so an attribute of a node that holds it is in the node's state by itself: state_dict returns the state with tensors that share memory with the relaxer's, and load_state_dict copies the tensors of the state into those. Only step_ moves r.

step_ abstractmethod

step_(r: Tensor, force: Tensor, dt: float, tau_h: float) -> None

Move r in place by dt seconds of the flow along force, without recording gradients.

PARAMETER	DESCRIPTION
r	The deviation h - mu of the node, (*batch, *out). TYPE: Tensor
force	The force \(F = (1 + \tau_h / \tau_r)\,r + \lambda + g\) on it, of the shape of r; the plain flow is \(\tau_h \dot r = -F\). TYPE: Tensor
dt	The time to move by, in seconds (0 changes nothing). TYPE: float
tau_h	The time constant \(\tau_h\) of r, in seconds (TimeConstants.h). TYPE: float

reset_ abstractmethod

reset_(samples: Tensor) -> None

Forget the state of samples (batch bool), whose r the node has just set to zero: the node calls it when it first sees a sample whose reset is set, where what the relaxer remembers is of the sample before. A relaxer without memory does nothing.

state_dict abstractmethod

state_dict() -> Any

Return the state. Its tensors may share memory with the object's: the run copies the state it takes.

load_state_dict abstractmethod

load_state_dict(state: Any) -> None

Restore the state that state_dict returned, in place.

Euler dataclass

Euler()

Bases: Relaxer

The plain gradient flow of r, \(\tau_h \dot r = -F\), by one explicit Euler step of dt: r -= (dt / tau_h) * force.

It is the default. The bound relaxer has no state (an empty state_dict).

Example::

hidden = pc.Layer(module, act_spec=..., tau=tau, optimizer=...,
                  relaxer=pc.relax.Euler(), learn_interval=24e-3)

bind

bind(r: Tensor) -> BoundRelaxer

Return the relaxer bound to r.

mujou.pc.loss

The losses of each sample for an output node.

A loss takes the prediction y_hat and the target y and returns the loss of each sample, of the shape of the batch. The node sums them; it never averages over the batch, so that the gradient it sends to the previous node is the credit of each sample (specs/pc.md). Each class holds its batch for the functions of functional of the same name.

MSELoss

MSELoss(batch: Size = ())

The mean squared error of each sample: functional.mse with batch.

Example::

loss = pc.loss.MSELoss()  # one sample: y_hat (out,) -> ()
loss = pc.loss.MSELoss(64)  # y_hat (64, out) -> (64,)

PARAMETER	DESCRIPTION
batch	The batch that the shapes start with; () for one sample. TYPE: Size DEFAULT: ()

CrossEntropyLoss

CrossEntropyLoss(batch: Size = ())

The cross entropy of each sample: functional.cross_entropy with batch.

Example::

loss = pc.loss.CrossEntropyLoss()  # one sample: y_hat (C,), y () -> ()
loss = pc.loss.CrossEntropyLoss(64)  # y_hat (64, C), y (64,) -> (64,)

PARAMETER	DESCRIPTION
batch	The batch that the shapes start with; () for one sample. TYPE: Size DEFAULT: ()

mujou.pc.msg

The messages of mujou.pc: the activity a node sends to the next one, what the host writes to an input node and what an output node sends.

Sample

Bases: TypedDict

What the host writes to an input node: the input x ((*batch, *in)), and the y, mask and reset of Activity.

reset is True for the samples that do not continue the ones written before: every sample of shuffled batches, or the first frame of a new episode in a stream. The nodes act on it once, when they first see the write.

x instance-attribute

x: Tensor

y instance-attribute

y: Tensor

mask instance-attribute

mask: Tensor

reset instance-attribute

reset: Tensor

Activity

Bases: TypedDict

The activity a node sends to the next nodes, with the target of each of its samples.

The tensors start with the batch, batch: () for one sample, (B,) for B samples, or more dimensions (see activity_spec).

• h: the activity, (*batch, *out).
• y: the target of each sample, (*batch, *target).
• mask: batch bool, whether each sample has a target. A sample whose mask is False adds nothing to the loss, whatever its y is.
• reset: batch bool, whether each sample starts anew: a node that first sees the activity sets the state it holds for those samples to zero (r and lam of a Layer). A sample that continues the one before it, as the next frame of a stream does, keeps it.
• time: the Timestamp that Run.write returned for the input h was computed from; Timestamp(ns=0) before the host first writes.

y, mask, reset and time come from the Sample the host wrote, and every node passes them on with the h it computed from it, in the same message: the output node pairs each prediction with its own target, however fast the host writes.

h instance-attribute

h: Tensor

y instance-attribute

y: Tensor

mask instance-attribute

mask: Tensor

reset instance-attribute

reset: Tensor

time instance-attribute

time: Timestamp

Prediction

Bases: TypedDict

What an output node sends: the prediction y_hat ((*batch, *out)), and the y, mask and time of the activity it was computed from (see Activity).

y_hat instance-attribute

y_hat: Tensor

y instance-attribute

y: Tensor

mask instance-attribute

mask: Tensor

time instance-attribute

time: Timestamp

activity_spec

activity_spec(h: TensorSpec, y: TensorSpec, batch: Size = ()) -> DictSpec[Activity]

Return the spec of an Activity whose h is of the spec h, whose targets are of the spec y and whose samples are batch.

batch is the shape that h and y start with, and that of mask: () (the default) for one sample, B or (B,) for B samples, or more dimensions, such as (E, A) for A agents in each of E environments. The samples must be independent of each other (Layer); a dimension that a module mixes, such as the time of a sequence, belongs to h, not to the batch.

Two calls with equal arguments return equal specs. A node of your own that connects to the nodes of mujou.pc declares its ports with it, or with the act_spec of the node before.

Example::

act = pc.msg.activity_spec(TensorSpec((256,)), TensorSpec((), torch.int64))
act = pc.msg.activity_spec(
    TensorSpec((64, 256)), TensorSpec((64,), torch.int64), batch=64
)

RAISES	DESCRIPTION
ValueError	h or y does not start with batch.

mujou.pc.clock

When a node that learns updates its parameters: LearnClock and the Tick it returns for each update.

LearnClock

LearnClock(*, period: float, learn_interval: float | None, max_dt: float)

Bases: Stateful

Decides, update by update, when a node updates its parameters, from the time the host wrote the node's input.

The node calls tick(dt, time, now) once in each update, with the dt of the update, the Timestamp the input it received was written at, and the time now on the run's timer. The clock learns in the first update at or after time + k * learn_interval (k = 1, 2, ...), but never in the update that first sees an input: its gradient is still that of the input before. The nodes on a path read the one timer of the run, so they learn at about the same time, however deep they are. With learn_interval=None, it learns in every update, a new input or not; with math.inf, never.

A Tick with learn counts as learned: the next interval starts there, whether the node updates its parameters or skips them for want of a target. elapsed sums the capped dt since then; an input that changes before the clock learned from it drops that sum, as an input without a target does, or inputs that change faster than learn_interval would pile the time up into one large update. dt is capped at max_dt, and the time beyond it is dropped from elapsed too.

The clock starts as if it had seen no input (no time a host writes at): the first update sees a new input, also one written at time 0.

The state (state_dict()) holds the time of the input seen last ("seen_ns", None before any), the time from which the next update learns ("learn_at_ns"), "elapsed", and whether the clock has learned since the input changed ("learned_input"). An attribute of a node that holds a LearnClock is in the node's state by itself.

It is not tied to PC-ALM: any node that learns from inputs the host writes can use it. LearnClock.of makes the clock of a node of PC-ALM from its time constants.

Example::

self.clock = pc.clock.LearnClock(period=self.period, learn_interval=24e-3,
                           max_dt=2e-3)
# In on_update(dt):
tick = self.clock.tick(dt, activity["time"], self.timer.now())
# ... move by tick.dt; if tick.learn, learn weighted by tick.elapsed.

PARAMETER	DESCRIPTION
period	The period of the node, in seconds. TYPE: float
learn_interval	The time in seconds from the write of an input to the first update of the parameters, and between the updates while the input stays. None learns in every update, math.inf never. TYPE: float | None
max_dt	The longest dt an update counts, in seconds. TYPE: float

RAISES	DESCRIPTION
ValueError	learn_interval is neither None nor > 0, max_dt is not > 0, or period is longer than max_dt.

of classmethod

of(
    tau: TimeConstants,
    *,
    period: float,
    learn_interval: float | None,
    max_dt: float | None = None,
) -> Self

Return the clock of a node of PC-ALM with the time constants tau, as Layer and Output make theirs: max_dt defaults to tau.max_dt.

RAISES	DESCRIPTION
ValueError	As LearnClock raises.

tick

tick(dt: float, time: Timestamp, now: Timestamp) -> Tick

Count an update of dt seconds that received an input written at time, at now on the run's timer, and return what the update does. With learn, the clock has learned at now.

state_dict

state_dict() -> dict[str, Any]

Return the state: plain numbers, which torch.load reads with weights_only=True.

load_state_dict

load_state_dict(state: Mapping[str, Any]) -> None

Restore the state that state_dict returned.

Tick

Bases: NamedTuple

What LearnClock.tick tells one update of a node.

dt instance-attribute

dt: float

The dt of the update, capped at max_dt: the time to move by.

new instance-attribute

new: bool

Whether the input is new: its time is not the one seen last.

learn instance-attribute

learn: bool

Whether to update the parameters in this update.

elapsed instance-attribute

elapsed: float

The sum of the capped dt since the parameters were last updated, this update included: the weight of the update when learn is True.

mujou.pc.functional

The stateless computations of PC-ALM that the nodes of mujou.pc are made of, for writing nodes of your own (such as one with several inputs).

A node i predicts mu = f(h_in; theta) from the activity h_in of the previous node, and keeps the deviation r = h - mu of its activity h and the dual lam of the constraint h = mu. Its energy is E = 0.5 * ||r + lam||^2 (delta = r + lam).

• vjp gives the gradients of E with respect to h_in and theta (one vector-Jacobian product inside the node). The node sends the one of h_in to the previous node, which adds the values it receives into its g.
• loss_vjp does the same for an output node, whose energy is the loss.
• relax_ moves r and lam by dt seconds (r as a relaxer of mujou.pc.relax says: pc.relax.Euler or one of your own); an optimizer of mujou.pc.optim (pc.optim.SGD, pc.optim.AdamW) moves theta.
• mse and cross_entropy are losses of each sample, for loss_vjp (MSELoss and CrossEntropyLoss hold their batch).

Every g here is a gradient of an energy; the update subtracts it. Gradients are summed over the batch, and the 1 / B (B the number of samples) is applied to theta only (the node divides g_theta by B before it gives it to the optimizer), so lam holds the credit of each sample. The shapes start with the batch, batch (see activity_spec).

vjp

vjp(
    module: Callable[[Tensor], Tensor],
    h_in: Tensor,
    delta: Callable[[Tensor], Tensor],
    params: Sequence[Tensor] = (),
) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]

Predict mu = module(h_in) and the gradients of the node's energy.

With the activity \(h\) held fixed, the energy \(E = \tfrac12\|h - \mu + \lambda\|^2\) has the gradients

\[ g_{\mathrm{in}} = \frac{\partial E}{\partial h_{\mathrm{in}}} = -J_h^\top \delta,\qquad g_\theta = \frac{\partial E}{\partial \theta} = -J_\theta^\top \delta, \]

where \(\delta = h - \mu + \lambda\) and \(J\) are the Jacobians of module at h_in. Both are summed over the batch. h_in is detached first, so no graph reaches the node that sent it.

PARAMETER	DESCRIPTION
module	The function f of the node, (*batch, *in) -> (*batch, *out). It must treat the samples independently. TYPE: Callable[[Tensor], Tensor]
h_in	The activity of the previous node, (*batch, *in). TYPE: Tensor
delta	Makes \(\delta\) of the node, (*batch, *out), from mu (detached, without gradients), such as lambda mu: r + lam. TYPE: Callable[[Tensor], Tensor]
params	The parameters to get the gradients of, usually those of module that require gradients. TYPE: Sequence[Tensor] DEFAULT: ()

RETURNS	DESCRIPTION
Tensor	(mu, g_in, g_theta): mu ((*batch, *out), detached), g_in
Tensor	((*batch, *in)) and g_theta (one tensor per parameter, in the order of
list[Tensor]	params; a parameter that module does not use gets zeros).

RAISES	DESCRIPTION
ValueError	delta does not make a tensor of the shape of mu.

loss_vjp

loss_vjp(
    module: Callable[[Tensor], Tensor],
    h_in: Tensor,
    y: Tensor,
    mask: Tensor,
    loss: Callable[[Tensor, Tensor], Tensor],
    params: Sequence[Tensor] = (),
) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]

Predict y_hat = module(h_in) and the gradients of the masked loss.

The energy of an output node is the loss of the samples that have a target,

\[ \ell = \sum_b \mathrm{mask}_b\, \ell_b(\hat y_b, y_b),\qquad g_{\mathrm{in}} = \frac{\partial \ell}{\partial h_{\mathrm{in}}},\qquad g_\theta = \frac{\partial \ell}{\partial \theta}. \]

A sample whose mask is False adds nothing, whatever its y is (NaN too): its y is replaced by zeros before loss sees it, and its g_in is \(J^\top 0\), zero as long as its y_hat is finite (NaN where it is not). With every mask False, g_in is thus zeros. h_in is detached first.

PARAMETER	DESCRIPTION
module	The function f of the node, (*batch, *in) -> (*batch, *out). TYPE: Callable[[Tensor], Tensor]
h_in	The activity of the previous node, (*batch, *in). TYPE: Tensor
y	The target, (*batch, *target). TYPE: Tensor
mask	Whether each sample has a target, batch bool. TYPE: Tensor
loss	The loss of each sample, (y_hat, y) -> batch, such as MSELoss(batch) or CrossEntropyLoss(batch). TYPE: Callable[[Tensor, Tensor], Tensor]
params	The parameters to get the gradients of. TYPE: Sequence[Tensor] DEFAULT: ()

RETURNS	DESCRIPTION
tuple[Tensor, Tensor, list[Tensor]]	(y_hat, g_in, g_theta) as vjp returns (mu, g_in, g_theta).

RAISES	DESCRIPTION
ValueError	loss does not return the shape of mask, as a loss that averages over the batch (nn.MSELoss()) does.

relax_

relax_(
    r: Tensor,
    lam: Tensor,
    g: Tensor,
    dt: float,
    tau: TimeConstants,
    relaxer: BoundRelaxer | None = None,
) -> None

Move r and lam in place by one explicit Euler step of dt seconds of

\[ \tau_h \dot r = -(r + \lambda + g) - \frac{\tau_h}{\tau_r}\,r,\qquad \tau_\lambda \dot\lambda = r. \]

r = h - mu is the deviation of the node's activity from its prediction, which stays in the node: the node sends mu. The step moves r first and then lam with the moved r. g is held fixed. The step is not split: g comes from the next nodes and stays as it was over dt, so the coupling between the nodes is explicit however finely the step is split. Keep dt short instead (Layer caps it at max_dt). tau.lam = inf leaves lam exactly as it is, and tau.r = inf adds no decay.

With a relaxer (pc.relax.Euler().bind(r) or one of your own), r moves as it says along the force \(F = (1 + \tau_h / \tau_r)\,r + \lambda + g\) instead of by the plain step: relaxer.step_(r, F, dt, tau.h). \(\lambda\) follows the moved r as before. The relaxer is the node's: it holds the state, and the node calls its reset_ where the force jumps.

PARAMETER	DESCRIPTION
r	The deviation h - mu of the node, (*batch, *out). TYPE: Tensor
lam	The dual of the node's activity, (*batch, *out). TYPE: Tensor
g	The sum of the gradients the next nodes sent, (*batch, *out). TYPE: Tensor
dt	The time to move by, in seconds. TYPE: float
tau	The time constants of r and lam. TYPE: TimeConstants
relaxer	How r moves, bound to r; None is the plain step. TYPE: BoundRelaxer | None DEFAULT: None

mse

mse(y_hat: Tensor, y: Tensor, batch: Size = ()) -> torch.Tensor

Return the mean squared error of each sample of batch.

The squared error is averaged over the dimensions of a sample, as nn.MSELoss averages it, but not over the batch: the mean of the result is nn.MSELoss()(y_hat, y), and with the 1 / B the node applies before its optimizer, theta follows its gradient.

With d values in a sample, its gradient with respect to y_hat is (2 / d) * (y_hat - y): that of 0.5 * ||y - y_hat||^2 for d = 2, twice it for d = 1. The larger it is, the more strongly the output pulls on the activities, and the shorter the period that keeps their explicit steps stable; for a scalar target, a loss of your own such as 0.5 * mse(y_hat, y, batch) keeps the pull of the half.

PARAMETER	DESCRIPTION
y_hat	The prediction, (*batch, *out). TYPE: Tensor
y	The target, of the same shape as y_hat. TYPE: Tensor
batch	The batch that the shapes start with; () for one sample. TYPE: Size DEFAULT: ()

RETURNS	DESCRIPTION
Tensor	The loss of each sample, batch.

RAISES	DESCRIPTION
ValueError	The shapes differ (a broadcast would silently give another shape), or they do not start with batch.

cross_entropy

cross_entropy(y_hat: Tensor, y: Tensor, batch: Size = ()) -> torch.Tensor

Return the cross entropy of each sample of batch.

The classes are the dimension right after the batch, as in F.cross_entropy ((N, C) and (C,) alike). A sample with more dimensions, such as one class for each pixel ((*batch, C, H, W)), has its cross entropy averaged over them, as nn.CrossEntropyLoss does, but not over the batch: the mean of the result is nn.CrossEntropyLoss()(y_hat, y). A class of -100 (the ignore_index of torch) adds zero but still counts in that mean; leave a sample out with mask instead.

PARAMETER	DESCRIPTION
y_hat	The logits, (*batch, C, *rest). TYPE: Tensor
y	The class of each sample, (*batch, *rest) int64, or the probability of each class, of the shape of y_hat. TYPE: Tensor
batch	The batch that the shapes start with; () for one sample. TYPE: Size DEFAULT: ()

RETURNS	DESCRIPTION
Tensor	The loss of each sample, batch.

RAISES	DESCRIPTION
ValueError	y_hat or y does not start with batch.