🔌 Ports and Channels¶

A port is a connection point of a node, and connect joins two ports with one channel. This page covers the kinds of ports, their specs, connecting, branching and merging, and what a port holds before anyone sends.

The kinds of ports¶

Kind	Peers	`recv()` returns
`Port`	At most one if it receives; any number if it only sends	The value
`MultiPort`	Any number, of one kind	A `dict[Channel, value]`, in `connect` order

Declare ports as attributes of the node in __init__, after super().__init__(); a run reads the attributes of the instance when it starts. A run raises TypeError for a port kept in a list, a tuple or a dict, and the class statement for a port declared as a class attribute: assign each port to an attribute of its own, or use a MultiPort.

Inputs of different roles, such as an input and its teacher signal, or a prediction and an error, go to separate Ports, each with a name that says its role. A MultiPort is for peers of one kind that the node combines the same way.

Specs¶

A port takes the spec of what it sends as send= and of what it receives as recv=, as keyword arguments; leave out the one a one-way port does not use.

self.prev = Port(recv=TensorSpec((32, 784)))  # receives only
self.next = MultiPort(send=TensorSpec((32, 256)), recv=TensorSpec((32, 256)))  # both ways

TensorSpec(shape, dtype) is one tensor; its dtype is torch.float32 unless given, whatever torch.get_default_dtype() is. A value sent must match the spec in shape and dtype: no dtype is converted. The dtypes a spec takes are in TensorSpec.dtypes. A spec also makes values of its shape and dtype, for the state of a node or a test: spec.zeros() and spec.empty() (a dict of them for a DictSpec), and TensorSpec.rand(generator=...).

A count, a flag or a time travels as a Python scalar, not as a tensor, so the receiver reads it without int(...):

Spec	Value	Initial value
`IntSpec(bits=64, signed=True)`	`int` of 8, 16, 32 or 64 bits	`0`
`FloatSpec()`	`float` (64 bits)	`0.0`
`BoolSpec()`	`bool`	`False`
`TimestampSpec()`	`mujou.time.Timestamp`	`Timestamp(ns=0)`

A message has a fixed size, so an IntSpec has a fixed number of bits: an int out of its range raises OverflowError when it is sent, and is never wrapped. A scalar spec has no shape and makes no values: write 0 or False.

Values on two ports travel independently: a reader may see a new value on one port with an old one on the other. Send values that belong together, such as an input and its target, as one message: a DictSpec gives a spec for each key of a TypedDict, the whole dict arrives at once, and type checkers know the type of each key. To use the spec of one key, take it with its kind: SAMPLE.field("x", TensorSpec).shape (TypeError if it is of another kind).

class Sample(TypedDict):
    x: torch.Tensor
    y: torch.Tensor


SAMPLE = DictSpec(Sample, x=TensorSpec((32, 4)), y=TensorSpec((32, 2)))

self.next = Port(send=SAMPLE)  # in the sender: self.next.send({"x": x, "y": y})
self.prev = Port(recv=SAMPLE)  # in the receiver: sample = self.prev.recv(); sample["x"]

Type checkers infer the type of a port from its specs. To annotate one (an argument of a helper, say), write Port[S, R] or MultiPort[S, R], where S is the type it sends and R the type it receives; a direction left out is Never, and a type checker rejects a send on a port that sends nothing. A function that takes any spec takes a MessageSpec[T].

from typing import Never

from mujou import MessageSpec, MultiPort, Port


def mean_of(port: MultiPort[Never, torch.Tensor]) -> torch.Tensor:
    return torch.stack(list(port.recv().values())).mean(0)


def sender[T](spec: MessageSpec[T]) -> Port[T]:
    return Port(send=spec)

Connecting¶

a.next.connect(b.prev) joins two ports and returns the Channel. Nodes and networks have no connect of their own: a node may have several inputs and outputs, and naming the two ports says which path you mean.

The spec that each port sends must match the spec the other receives, in both directions; otherwise connect raises TypeError
A Port that receives takes one peer: connecting a second raises ValueError. Receive from several with a MultiPort
Connecting the same two ports again returns the channel they have. A port may connect to a port of its own node
Connections may form cycles: a node may receive from a node it sends to, as the nodes of mujou.pc do on two-way ports
A port connects only to ports of nodes in the same run; mujou.run raises ValueError for a peer outside the tree

Fan-out and merging¶

A Port that only sends takes any number of peers, and sends every value to each of them. A MultiPort merges: its recv() returns the latest value of each channel. The value of a peer that has never sent is the initial value of the peer's port.

class Apply(Node):
    """Sends ``fn(x)``."""

    def __init__(self, fn: Callable[[torch.Tensor], torch.Tensor]) -> None:
        super().__init__()
        self.fn = fn
        self.prev = Port(recv=VECTOR)
        self.next = Port(send=VECTOR)  # only sends: any number of peers

    @override
    def on_update(self, dt: float) -> None:
        super().on_update(dt)
        self.next.send(self.fn(self.prev.recv()))


class Sum(Node):
    """Sends the sum of the latest messages of all its peers."""

    def __init__(self) -> None:
        super().__init__()
        self.prev = MultiPort(recv=VECTOR)  # merges any number of peers
        self.next = Port(send=VECTOR)

    @override
    def on_update(self, dt: float) -> None:
        super().on_update(dt)
        received = self.prev.recv()  # dict[Channel, torch.Tensor], in connect order
        self.next.send(torch.stack(list(received.values())).sum(dim=0))


class Diamond(mujou.Network):
    """``shift -> (double, square) -> sum``: y = 2 (x + 1) + (x + 1)^2."""

    def __init__(self) -> None:
        super().__init__()
        self.shift = Apply(lambda x: x + 1)
        self.double = Apply(lambda x: 2 * x)
        self.square = Apply(lambda x: x**2)
        self.sum = Sum()
        self.shift.next.connect(self.double.prev)  # one port, two peers
        self.shift.next.connect(self.square.prev)
        self.double.next.connect(self.sum.prev)  # two peers, one MultiPort
        self.square.next.connect(self.sum.prev)
        self.prev, self.next = self.shift.prev, self.sum.next

The receiving node decides how to combine the values: sum them, average them, weight each peer. To treat the peers differently, keep the Channel that connect returned for each, as an attribute of the receiving node (self.sum.from_double = self.double.next.connect(self.sum.prev)), and index with it: self.prev.recv()[self.from_double].

Sending to one channel¶

A MultiPort has two ways to send:

How	Where it goes	What `Run.read` returns
`send(value)`	The same value to every peer	This value
`send_to(channel, value)`	Only to the peer at the end of `channel`	Unchanged

class Split(Node):
    """Sends ``x`` on the channel ``plus`` and ``-x`` on the channel ``minus``."""

    # Channels of self.next, set by the network to the ones connect returned.
    plus: Channel
    minus: Channel

    def __init__(self) -> None:
        super().__init__()
        self.prev = Port(recv=VECTOR)
        self.next = MultiPort(send=VECTOR)  # sends to each peer on its own

    @override
    def on_update(self, dt: float) -> None:
        super().on_update(dt)
        x = self.prev.recv()
        self.next.send_to(self.plus, x)
        self.next.send_to(self.minus, -x)


class Magnitude(mujou.Network):
    """``split -> (relu, relu) -> sum``: y = relu(x) + relu(-x) = |x|."""

    def __init__(self) -> None:
        super().__init__()
        self.split = Split()
        self.pos, self.neg = Apply(torch.relu), Apply(torch.relu)
        self.sum = Sum()
        # Hand split the channels it sends on: the order of connect does not matter.
        self.split.plus = self.split.next.connect(self.pos.prev)
        self.split.minus = self.split.next.connect(self.neg.prev)
        self.pos.next.connect(self.sum.prev)
        self.neg.next.connect(self.sum.prev)
        self.prev, self.next = self.split.prev, self.sum.next

The network hands the node the channels connect returned (self.split.plus = ...). Do not pick channels by their order in self.next.channels instead: swapping two connect calls would swap the destinations without any error. An attribute that holds a Channel is not part of the state.

The host reads with Run.read only what a port sent with send (its initial value until then); it cannot read the values of send_to, nor a port that only receives.

Before anyone sends¶

Until a port first sends, its peers, and the host that reads it, get its initial value: zeros (the initial values above for scalars), or the value of initial=, which must match send and is copied when the port is made. Its time is None.

self.next = Port(send=TensorSpec((32, 256)), initial=torch.full((32, 256), 0.5))

recv_with_time() returns the value with the time it was sent (a mujou.time.Timestamp), or None before the first send. Compare it with the time of the value read last to tell a new value from the same one read again. A port that the host writes counts as sent from the inputs= of the run, or from the host's first write: see Host I/O.

Exposed ports¶

A network exposes a child's port by assigning it to an attribute of its own (self.prev = self.l1.prev). The attribute is only a reference: the port still belongs to the node, and the outside connects to it as to the node's (x.next.connect(net.prev)). By convention, a network exposes the port of its first node as prev and that of its last as next, so that it connects like a node.