Nodes and Ports¶
mujou.Node ¶
Node(
*,
period: float = 0.01,
seed: int | None = None,
device: Device = "cpu",
dedicated_thread: bool = False,
)
Bases: Component
A unit of asynchronous computation: the leaf of a network.
Call super().__init__(...) first in __init__ (mujou.run raises
TypeError for a node that did not), then declare the node's ports and
state as attributes of the instance. A run reads the attributes when it
starts. A node must not hold other nodes or networks as attributes, nor
ports in a list, a tuple or a dict (the run raises TypeError); put
nodes side by side in a Network instead.
Override the lifecycle hooks below; each does nothing by default. As in any
subclass, call the same hook of super() first in each override. The
framework calls the hooks on the thread that runs the node (see "Threads"
below), in this order:
on_start(): once, when the run starts (also when it is restored). Open the node's I/O here. A node withdedicated_thread=Truegets it on its own thread, so it can also make its per-thread settings here. Do not send anything that depends on the state here: a restore loads the state after it.load_state_dict(state): only when the run was restored, afteron_startand before the firston_update.on_update(dt): repeatedly, once everyperiodon average (see below).dtis the time in seconds on the run's timer since the node's last update. Read the ports withrecv, update in proportion todt, and write the ports withsend. It never waits for other nodes.on_pause(): when the node has stopped for the run to take its state (Run.state_dict).on_resume(): when that pause ends, before the nexton_update. A pause of the timer (Timer.pause) calls neither: the node only sees no time pass.on_exception(): once on every node, when a hook of any node has raised, beforeon_stop.on_stop(): once, when the run stops, also after an exception. It must not change whatstate_dictreturns: the state taken after a normal end is the state at the stop.
state_dict() is called when the run's state is taken, while the node is
paused (or, after a normal end, on the host's thread that called
Run.state_dict). Its ports refuse to receive or send there, so that
taking the state is not seen by the other nodes. The framework copies what
it returns before the node resumes.
The state of a node is, as for an nn.Module, every attribute that holds
a torch.Tensor (nn.Parameter included) or a Stateful (an
nn.Module, a pc.clock.LearnClock; not a class), every attribute
registered with register_state, and the state of generator. What an
attribute holds when the state is taken or loaded decides. The settings that __init__
computes from its arguments (a time constant, say) are no state: a node
restored with other arguments keeps its own. Register the plain values that
change while the node runs (a counter, the time last seen), and the
containers of tensors (a list as a replay buffer, a dict of tensors): a
list or a dict is no state by itself, even of tensors. Make the state
in __init__ and update it in place: load_state_dict copies into the
tensors that the node holds. To choose the state yourself, override both
state_dict and load_state_dict: extend what super().state_dict()
returns, and pass super().load_state_dict the state without the keys
you added.
If any hook raises, the whole run stops. A hook must not call mujou.run
or the methods of any Run (write, read, state_dict, ...):
they raise RuntimeError there, because a node talks only through its own
ports.
Threads: by default the node runs on the workers, a pool of threads that
the run shares among its nodes (mujou.run(..., workers=...)). Each hook
may then run on a different thread from the previous one, but the hooks of
a node never run at the same time, and each sees what the previous one
wrote. Whether the next hook sees what a hook set per thread
(threading.local, threading.current_thread(), settings of torch
such as torch.set_grad_enabled or torch.autocast) depends on the
schedule, and the next node on the same worker may see it too: set such
settings back before the hook returns (inside a with, say).
contextvars do persist: each node runs its hooks in a
contextvars.Context of its own, which starts empty (it does not inherit
the host's). A hook that blocks (time.sleep, waiting for I/O or a
lock) holds its worker meanwhile, and the other nodes' dt grows; give
such a node dedicated_thread=True, so that it runs on a thread of its
own, with every hook on that thread.
The framework leaves gradient recording as it is, on by default as in
plain PyTorch: a hook may call backward directly, and computes without
a graph inside with torch.no_grad():. No graph crosses nodes: a port
sends a detached copy. Before each hook, when device is CUDA, the
framework makes the device and a stream of the node's own current: the
only per-thread settings it makes, and around state_dict (which may
run on the host's thread) it sets the thread's own back. On CUDA it also
waits for the work queued on that stream after on_update (holding its
worker meanwhile), after load_state_dict and around state_dict. It
does this around the hooks, so overriding one never skips it.
The node updates once every period (0.01 s by default) of the run's
timer on average: the time an update takes counts toward the period, and a
node that falls a whole period behind does not update several times in
a row to catch up. A pause does not count. dt is still the time that
actually passed since the last update, longer than period when the
node is late, so keep updating in proportion to it and stay stable for the
longest dt it can get. A manual run (ManualRun.step) ignores the
period.
period=0.0 updates as soon as any time has passed since the last update.
Set it only when the node must: such a node keeps a CPU core busy, and takes
time from the other nodes, whose dt it makes irregular.
period=math.inf never updates by itself: only ManualRun.step updates
such a node.
Example::
class Hidden(Node):
def __init__(self, n_in: int, n_out: int) -> None:
super().__init__(seed=0)
x, h = TensorSpec((32, n_in)), TensorSpec((32, n_out))
# Send W^T eps back and receive x from the previous side (one peer).
self.prev = Port(send=x, recv=x)
# Send h forward and receive W^T eps from each next peer (many).
self.next = MultiPort(send=h, recv=h)
@override
def on_update(self, dt: float) -> None:
super().on_update(dt)
x = self.prev.recv() # torch.Tensor
feedback = self.next.recv() # dict[Channel, torch.Tensor]
...
self.next.send(h)
self.prev.send(wt_eps)
l1, l2 = Hidden(784, 256), Hidden(256, 256)
l1.next.connect(l2.prev)
| PARAMETER | DESCRIPTION |
|---|---|
period
|
The time between two updates on average, in seconds of the
run's timer.
TYPE:
|
seed
|
The seed of
TYPE:
|
device
|
Where the node computes: the CPU or a CUDA device (no other
device type is supported). Received messages arrive on it.
TYPE:
|
dedicated_thread
|
Run the node on a thread of its own instead of the
run's workers: for a node whose hooks block, or that relies on
per-thread state set once (in
TYPE:
|
| RAISES | DESCRIPTION |
|---|---|
ValueError
|
|
device
property
¶
The device the node computes on and receives messages on.
Allocate with it (torch.zeros(..., device=self.device)): the current
device in the hooks is this one, but an explicit device states the
intent.
period
property
¶
The time between two updates on average, in seconds of the run's
timer. 0.0 updates as soon as any time has passed.
dedicated_thread
property
¶
Whether the node runs on a thread of its own instead of the run's workers.
generator
property
¶
The node's own random generator on device, seeded with seed.
seed sets nothing else: use this generator instead of the global
one, which all nodes share (do not call torch.manual_seed in a
node), also to initialize modules (nn.init.normal_(w,
generator=self.generator)). super().__init__ creates it, so it can
be used in __init__ already, and the hooks go on with the same
sequence. Its state is part of the node's state, so a restored node
goes on with the same sequence too.
timer
property
¶
The reading side of the timer of the run whose hook runs, which
measures the dt of its nodes.
Read the time with self.timer.now() in the hooks (state_dict
included). The timer is the one given to mujou.run(..., timer=...),
the global timer by default, or, for a run restored from a state, a
new timer that started at the time of the state (Run.timer). A
ReadOnlyTimer only reads: a node does not change the time of its
run (Timer.sleep, set_time_scale and VirtualTimer.advance
raise RuntimeError in a hook). It stops at the time the run ends,
so the hooks that run after the end (on_exception, on_stop)
see the time stand still; measure real time there instead. The host
reads Run.timer.
| RAISES | DESCRIPTION |
|---|---|
RuntimeError
|
This is not read in a hook of this node while it runs
(but in |
register_state ¶
Add the attribute name to the state, whatever value it holds.
Use it for the plain values that change while the node runs (a
counter, the time last seen) and for the containers of tensors (a list
as a replay buffer), which are not state by themselves. The
attribute is state while it is assigned, whatever it holds. The default
state_dict returns the value as it is, and the default
load_state_dict assigns a deep copy of the saved value (or copies
into a tensor, or loads into a Stateful, that the attribute holds
then).
Register in __init__: a restore loads into new instances, which
must have the same state.
The state is data: when the run takes it, it detaches every tensor of
what state_dict returns (in its dicts, lists and tuples) and copies
the result, so no graph goes with it.
torch.load(..., weights_only=True) reads only basic types (int,
float, bool, str, None, and lists, tuples and dicts of
them and of tensors): keep the registered values to those if the state
is saved with torch.save.
| RAISES | DESCRIPTION |
|---|---|
AttributeError
|
The node has no attribute |
on_start ¶
Called once when the run starts (also when it is restored), before any other hook. Do not send anything that depends on the state here: a restore loads it later.
on_update ¶
Called repeatedly with dt, the seconds on the run's timer since the
node's last update (always > 0). Read the ports, update in proportion to
dt and send.
on_pause ¶
Called when the node has stopped for the run to take its state
(Run.state_dict), at the end of its current on_update. A pause
of the timer (Timer.pause) does not call it.
on_resume ¶
Called when the pause for taking the state ends, before the next
on_update.
on_exception ¶
Called once on every node when a hook of any node has raised, before
on_stop. What it raises is logged and ignored.
on_stop ¶
Called once when the run stops, also after an exception. It must not
change what state_dict returns. What it raises is logged and
ignored.
state_dict ¶
Return the node's state: "generator", the state of
generator, then each attribute of the state in the order of
assignment, under its name: a tensor detach-ed, the state_dict()
of a Stateful, or a registered value as it is.
The tensors share memory with the attributes (the framework copies the
tree). requires_grad and .grad are not part of the state, as in
nn.Module.state_dict.
load_state_dict ¶
Restore state, as the default state_dict returned it, into the
attributes of the state, without recording gradients: copy_ into
each tensor, load_state_dict into each Stateful, and
assign a copy.deepcopy of the saved value to any other registered
attribute, with the tensors in it (in its dicts, lists and tuples)
moved to device.
The objects stay the same, so references taken in __init__ or
on_start see the restored values, and a tensor that records
gradients stays a leaf that does. copy_ casts the dtype and moves
the device. The state of generator is set as it is saved, and a
generator keeps its state in a format of its device's type: the state
must come from a node on the same type of device (the CPU or CUDA).
An override that keeps a tensor of state as it is shares its memory
with the tree given to mujou.run(..., state=...): copy it instead.
| RAISES | DESCRIPTION |
|---|---|
ValueError
|
The keys of |
mujou.Stateful ¶
Bases: ABC
An object with a state of its own, which a node keeps in its state.
An attribute of a Node that holds a Stateful is in the node's state
by itself: the default Node.state_dict returns its state_dict()
under the attribute's name, and Node.load_state_dict passes that back
to its load_state_dict. An object is one by inheriting this class, not
by having methods of these names. nn.Module is registered as one;
register a class of another library with Stateful.register(cls).
Example::
class Counter(mujou.Stateful):
def __init__(self) -> None:
self.n = 0
def state_dict(self) -> dict[str, int]:
return {"n": self.n}
def load_state_dict(self, state: Mapping[str, int]) -> None:
self.n = state["n"]
mujou.Port ¶
Port(
*,
send: MessageSpec[S] | None = None,
recv: MessageSpec[R] | None = None,
initial: S | None = None,
)
Bases: _Port[S, R]
A connection point that sends messages of type S and receives at most
one peer's messages of type R.
Declare ports as attributes of a node in __init__. send is the spec of
what the port sends and recv the spec of what it receives; omit one for a
one-way port (its message type is then Never). recv() returns one
message, so a port that receives takes one peer; a port that only sends takes
any number, and send reaches them all. To receive from several peers of
one kind, use a MultiPort; for inputs of different roles, use one port
each.
Two ports connect when the send of each equals the recv of the other.
A port may stay unconnected: the host then writes what it receives
(Run.write).
initial is what the port has sent before its first send: what its
peers receive and the host reads until then. A run copies it when it
starts. Without it, that is zeros. A restored run
takes the values from the state instead.
A run names the port in its errors by the path of its node and its
attribute ("body.l1.next"; "Shift.next" for a port of the node run
as the root).
Example::
class Sensor(Node):
def __init__(self) -> None:
super().__init__()
self.next = Port(send=TensorSpec((3,)), initial=torch.ones(3))
class Monitor(Node):
def __init__(self) -> None:
super().__init__()
self.prev = Port(recv=TensorSpec((3,)))
@override
def on_update(self, dt: float) -> None:
super().on_update(dt)
x = self.prev.recv() # torch.Tensor; ones until the sensor sends
sensor, monitor = Sensor(), Monitor()
sensor.next.connect(monitor.prev)
send_spec
property
¶
The spec of the messages this port sends, or None.
recv_spec
property
¶
The spec of the messages this port receives, or None.
initial
property
¶
What the port has sent before its first send, or None for
zeros.
channels
property
¶
The channels of this port, in the order they were connected.
connect ¶
Connect this port to other and return the channel between them.
The channel is registered at both ports, so each sends to and receives from the other. Connecting the same two ports again, in either order, returns the same channel. A port may connect to itself.
| RAISES | DESCRIPTION |
|---|---|
TypeError
|
The specs are not mutually reversed: |
ValueError
|
A |
send ¶
Send msg to every peer and to what the host reads (Run.read).
msg must match the send spec exactly; tensor dtypes are never
converted (a scalar is taken as Python's typing takes it: an int
for a float). The
values are detached and copied, so changing msg afterwards changes
nothing that was sent. If a node on CUDA reads the port, a tensor on
the node's CUDA device is copied on the device, on the node's stream,
and send waits for that stream before the peers can read it; a
peer on CUDA copies it from device to device. It stays on the device
also while the hook has made another device current
(torch.cuda.device(...)). Any other CUDA tensor is copied to the
host. It never waits for the peers to read.
| RAISES | DESCRIPTION |
|---|---|
RuntimeError
|
The caller is not a hook of the node that owns the port,
while it runs (or is its |
TypeError
|
A value is not of its field's kind (a |
ValueError
|
A shape or a dtype differs from the spec. |
OverflowError
|
An |
BufferError
|
A tensor is not dense ( |
recv ¶
Return the latest message received, without waiting for a new one.
The message comes from the peer's send; an unconnected port reads
what the host writes with Run.write instead. Until that writer first
writes, it is the initial value: the peer port's initial, or else
zeros. For a port the host writes, it is the value in inputs of
mujou.run, or else zeros. A restored run takes
it from state. recv_with_time tells whether the writer has
written yet.
Every call copies into new tensors on the owning node's device, without gradients, so changing them in place changes nothing else.
| RAISES | DESCRIPTION |
|---|---|
RuntimeError
|
The caller is not a hook of the node that owns the port,
while it runs (or is its |
recv_with_time ¶
Like recv, but return (message, time).
time is the Timestamp on the run's timer at which the writer wrote
the message, or None while the writer has not written yet (the
message is then the initial value). The writer is the peer (its send
or send_to), or the host for an unconnected port (inputs of
mujou.run or Run.write). A restored run keeps the times of the
state, None included. Raises as recv does.
mujou.MultiPort ¶
MultiPort(
*,
send: MessageSpec[S] | None = None,
recv: MessageSpec[R] | None = None,
initial: S | None = None,
)
Bases: _Port[S, R]
A connection point that sends messages of type S and receives
messages of type R from any number of peers of one kind.
It is declared and connected as a Port is, but recv() returns
{channel: message}, one entry per peer in connection order, and
send_to sends to one peer. Use it to merge or to branch peers that play
the same role; for inputs of different roles (a prediction and an error,
say), use one Port each. The host does not write to a MultiPort.
Example::
class Sum(Node):
def __init__(self) -> None:
super().__init__()
self.prev = MultiPort(recv=TensorSpec((3,)))
self.next = Port(send=TensorSpec((3,)))
@override
def on_update(self, dt: float) -> None:
super().on_update(dt)
self.next.send(sum(self.prev.recv().values(), torch.zeros(3)))
send_spec
property
¶
The spec of the messages this port sends, or None.
recv_spec
property
¶
The spec of the messages this port receives, or None.
initial
property
¶
What the port has sent before its first send, or None for
zeros.
channels
property
¶
The channels of this port, in the order they were connected.
connect ¶
Connect this port to other and return the channel between them.
The channel is registered at both ports, so each sends to and receives from the other. Connecting the same two ports again, in either order, returns the same channel. A port may connect to itself.
| RAISES | DESCRIPTION |
|---|---|
TypeError
|
The specs are not mutually reversed: |
ValueError
|
A |
send ¶
Send msg to every peer and to what the host reads (Run.read).
msg must match the send spec exactly; tensor dtypes are never
converted (a scalar is taken as Python's typing takes it: an int
for a float). The
values are detached and copied, so changing msg afterwards changes
nothing that was sent. If a node on CUDA reads the port, a tensor on
the node's CUDA device is copied on the device, on the node's stream,
and send waits for that stream before the peers can read it; a
peer on CUDA copies it from device to device. It stays on the device
also while the hook has made another device current
(torch.cuda.device(...)). Any other CUDA tensor is copied to the
host. It never waits for the peers to read.
| RAISES | DESCRIPTION |
|---|---|
RuntimeError
|
The caller is not a hook of the node that owns the port,
while it runs (or is its |
TypeError
|
A value is not of its field's kind (a |
ValueError
|
A shape or a dtype differs from the spec. |
OverflowError
|
An |
BufferError
|
A tensor is not dense ( |
recv ¶
Return {channel: message}, the latest message of each peer in
connection order ({} while unconnected), as Port.recv returns the
one message of its peer.
| RAISES | DESCRIPTION |
|---|---|
RuntimeError
|
As |
recv_with_time ¶
Like recv, but each message comes as (message, time), as
Port.recv_with_time gives it.
send_to ¶
Send msg only to the peer across channel, one of channels
(a key of recv).
Unlike send, it leaves what the host reads (Run.read) unchanged:
before the first send, that is this port's initial (or
zeros). Copies and raises as send does.
mujou.Channel ¶
One connection between two ports, which connect makes.
A channel is the handle of a peer: a MultiPort's recv is keyed by
channel, and send_to takes one. A channel made otherwise connects
nothing.
Example::
ch = a.next.connect(b.prev)
ch.other(a.next) is b.prev # True
other ¶
Return the end of this channel that is not port (port itself
for a port connected to itself).
mujou.MessageSpec ¶
Bases: Generic[_M]
The spec of a message of type _M.
TensorSpec is a MessageSpec[torch.Tensor], IntSpec a
MessageSpec[int] and DictSpec(Sample, ...) a MessageSpec[Sample].
Annotate with it a function that takes any spec::
def sender[T](spec: MessageSpec[T]) -> Port[T]:
return Port(send=spec)
mujou.TensorSpec ¶
Bases: MessageSpec[Tensor]
The shape and the dtype of a message that is one torch.Tensor.
Every dimension is at least 1, and () is a scalar. A batch dimension is
just a part of the shape.
Two specs are equal when their shapes and dtypes are equal. A message whose shape or dtype differs from the spec is an error; dtypes are never converted. A spec has no device: a message of the spec may be on any device.
zeros, rand and empty make a tensor of the spec, on the CPU unless
device says otherwise. For any other initial value, fill one of them in
place (spec.empty().fill_(0.5)).
Example::
spec = TensorSpec((32, 256)) # float32
spec = TensorSpec((32, 256), torch.bfloat16)
x0 = spec.zeros()
noise = TensorSpec((4,)).rand(generator=torch.Generator().manual_seed(0))
ones = spec.empty(device="cuda").fill_(1.0)
| RAISES | DESCRIPTION |
|---|---|
TypeError
|
|
zeros ¶
Return a tensor of this spec filled with zeros on device.
empty ¶
Return a tensor of this spec on device whose values are not set.
Set them before use, in place: spec.empty().fill_(0.5).
rand ¶
Return a tensor of this spec drawn uniformly from [0, 1) on device.
Pass a generator on device to make the result deterministic;
without one, the global torch generator is used.
mujou.IntSpec ¶
Bases: MessageSpec[int]
A message that is one Python int of bits bits, signed or not.
A value out of the range of bits and signed is an error when it is
sent or written (OverflowError); it is never wrapped. A bool is an
int and is taken as 0 or 1. The initial value is 0.
Two specs are equal when their bits and signed are equal.
Example::
count = IntSpec() # int64
level = IntSpec(8, signed=False) # 0 to 255
| RAISES | DESCRIPTION |
|---|---|
ValueError
|
|
mujou.FloatSpec ¶
Bases: _Unit, MessageSpec[float]
A message that is one Python float (64 bits).
An int is taken too, as a float is in Python's typing, and is read
back as a float. The initial value is 0.0. Every FloatSpec is equal.
mujou.BoolSpec ¶
Bases: _Unit, MessageSpec[bool]
A message that is one Python bool. The initial value is False. Every
BoolSpec is equal.
mujou.TimestampSpec ¶
Bases: _Unit, MessageSpec[Timestamp]
A message that is one mujou.time.Timestamp, such as a time that
Run.write returned, carried on with what was computed from it.
The initial value is Timestamp(ns=0). Every TimestampSpec is equal.
mujou.DictSpec ¶
Bases: MessageSpec[TD]
The spec of a message that is a TypedDict of fields.
Give the TypedDict class and one FieldSpec per key: a TensorSpec
for a key annotated as torch.Tensor, IntSpec for int,
FloatSpec for float, BoolSpec for bool and
TimestampSpec for Timestamp. Every key of the TypedDict is
required (no nesting), and the keyword arguments name exactly its keys. The
keys keep the order in which the TypedDict defines them.
Two specs are equal when they use the same TypedDict class and equal
specs for every key.
Example::
class Feedback(TypedDict):
error: torch.Tensor
count: int
spec = DictSpec(Feedback, error=TensorSpec((32, 256)), count=IntSpec())
msg: Feedback = spec.zeros()
error = spec.field("error", TensorSpec) # a TensorSpec, to type checkers too
| RAISES | DESCRIPTION |
|---|---|
TypeError
|
The keyword arguments miss a key of the |
specs
property
¶
The spec of each key, in the order the TypedDict defines.
To use the spec of one key as its kind, take it with field.
field ¶
Return the spec of key, checked to be a kind.
Example::
h = act_spec.field("h", TensorSpec)
h.shape
| RAISES | DESCRIPTION |
|---|---|
KeyError
|
No such key. |
TypeError
|
The spec of |
zeros ¶
Return a message whose tensors are zeros on device and whose
scalars are zero (0, 0.0, False, Timestamp(ns=0)).
empty ¶
Return a message whose tensors are on device and not set, and
whose scalars are zero as in zeros.
Set the tensors before use, in place: msg["error"].fill_(0.5).