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):

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;
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;
moves r and then lam by dt with the g the next nodes sent (functional.relax_);
sends h = mu, and the y, mask, reset and time that came with h_in;
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`.

PC-ALM¶

mujou.pc ¶

Chain ¶

input instance-attribute ¶

output instance-attribute ¶

prev instance-attribute ¶

next instance-attribute ¶

layers property ¶

Shown ¶

prev instance-attribute ¶

next instance-attribute ¶

Input ¶

act_spec instance-attribute ¶

prev instance-attribute ¶

next instance-attribute ¶

on_update ¶

Layer ¶

clock instance-attribute ¶

module instance-attribute ¶

optimizer instance-attribute ¶

r instance-attribute ¶

relaxer instance-attribute ¶

lam instance-attribute ¶

act_spec instance-attribute ¶

prev instance-attribute ¶

next instance-attribute ¶

on_update ¶

Output ¶

clock instance-attribute ¶

module instance-attribute ¶

optimizer instance-attribute ¶

prev instance-attribute ¶

next instance-attribute ¶

on_update ¶

TimeConstants dataclass ¶

h instance-attribute ¶

lam instance-attribute ¶

r instance-attribute ¶

period property ¶

max_dt property ¶

show ¶

mujou.pc.optim ¶

Optimizer ¶

bind abstractmethod ¶

BoundOptimizer ¶

step_ abstractmethod ¶

state_dict abstractmethod ¶

load_state_dict abstractmethod ¶

SGD dataclass ¶

tau instance-attribute ¶

bind ¶

AdamW dataclass ¶

tau instance-attribute ¶

tau_m instance-attribute ¶

tau_v instance-attribute ¶

eps class-attribute instance-attribute ¶

weight_decay class-attribute instance-attribute ¶

bind ¶

mujou.pc.relax ¶

Relaxer ¶

bind abstractmethod ¶

BoundRelaxer ¶

step_ abstractmethod ¶

reset_ abstractmethod ¶

state_dict abstractmethod ¶

load_state_dict abstractmethod ¶

Euler dataclass ¶

bind ¶

mujou.pc.loss ¶

MSELoss ¶

CrossEntropyLoss ¶

mujou.pc.msg ¶

Sample ¶

x instance-attribute ¶

y instance-attribute ¶

mask instance-attribute ¶

reset instance-attribute ¶

Activity ¶

h instance-attribute ¶

y instance-attribute ¶

input `instance-attribute` ¶

output `instance-attribute` ¶

prev `instance-attribute` ¶

next `instance-attribute` ¶

layers `property` ¶

prev `instance-attribute` ¶

next `instance-attribute` ¶

act_spec `instance-attribute` ¶

prev `instance-attribute` ¶

next `instance-attribute` ¶

clock `instance-attribute` ¶

module `instance-attribute` ¶

optimizer `instance-attribute` ¶

r `instance-attribute` ¶

relaxer `instance-attribute` ¶

lam `instance-attribute` ¶

act_spec `instance-attribute` ¶

prev `instance-attribute` ¶

next `instance-attribute` ¶

clock `instance-attribute` ¶

module `instance-attribute` ¶

optimizer `instance-attribute` ¶

prev `instance-attribute` ¶

next `instance-attribute` ¶

TimeConstants `dataclass` ¶

h `instance-attribute` ¶

lam `instance-attribute` ¶

r `instance-attribute` ¶

period `property` ¶

max_dt `property` ¶

bind `abstractmethod` ¶

step_ `abstractmethod` ¶

state_dict `abstractmethod` ¶

load_state_dict `abstractmethod` ¶

SGD `dataclass` ¶

tau `instance-attribute` ¶

AdamW `dataclass` ¶

tau `instance-attribute` ¶

tau_m `instance-attribute` ¶

tau_v `instance-attribute` ¶

eps `class-attribute` `instance-attribute` ¶

weight_decay `class-attribute` `instance-attribute` ¶

bind `abstractmethod` ¶

step_ `abstractmethod` ¶

reset_ `abstractmethod` ¶

state_dict `abstractmethod` ¶

load_state_dict `abstractmethod` ¶

Euler `dataclass` ¶

x `instance-attribute` ¶

y `instance-attribute` ¶

mask `instance-attribute` ¶

reset `instance-attribute` ¶

h `instance-attribute` ¶

y `instance-attribute` ¶

mask `instance-attribute` ¶

reset `instance-attribute` ¶

time `instance-attribute` ¶

y_hat `instance-attribute` ¶

y `instance-attribute` ¶

mask `instance-attribute` ¶

time `instance-attribute` ¶

of `classmethod` ¶

dt `instance-attribute` ¶

new `instance-attribute` ¶

learn `instance-attribute` ¶

elapsed `instance-attribute` ¶