🚀 1. Run a Network¶

This page runs the smallest program, two nodes that compute \(y = 2x + 1\), and explains it line by line. Then it branches and merges a network.

flowchart LR
    H1["Host"] -->|"writes x"| S["scale (×2)"]
    S --> F["shift (+1)"]
    F -->|"reads y"| H2["Host"]

The program¶

examples/quickstart.py

"""The smallest program: two nodes in a network, running in the background.

The host writes an input, waits until the output arrives, and stops.
"""

from typing import override

import torch

import mujou
from mujou import Network, Node, Port, TensorSpec

VECTOR = TensorSpec((3,))  # every message here is a float32 vector of 3


class Affine(Node):
    """Sends ``a * x + b``, where ``x`` is the latest input."""

    def __init__(self, a: float, b: float) -> None:
        super().__init__()
        self.a, self.b = a, b  # settings, not state
        self.prev = Port(recv=VECTOR)  # receives from the previous node
        self.next = Port(send=VECTOR)  # sends to the next node

    @override
    def on_update(self, dt: float) -> None:
        super().on_update(dt)
        self.next.send(self.a * self.prev.recv() + self.b)  # recv() never waits




class Pipeline(Network):
    """``x -> scale -> shift ->``: y = 2 x + 1."""

    def __init__(self) -> None:
        super().__init__()
        self.scale = Affine(2.0, 0.0)
        self.shift = Affine(1.0, 1.0)
        self.scale.next.connect(self.shift.prev)
        self.prev, self.next = self.scale.prev, self.shift.next  # exposed as its own




net = Pipeline()
x = torch.tensor([1.0, 2.0, 3.0])
# The nodes run on the run's worker threads until the end of the `with` block.
with mujou.run(net) as run:
    run.write({net.prev: x})
    # x takes a few updates to pass through the nodes: wait until y arrives, with
    # the time it was sent (TimeoutError after 5 s of real time).
    y, t = run.wait_for(net.next, lambda y: torch.equal(y, 2 * x + 1), timeout=5.0)
    assert t is not None  # None only for the zeros before the first send

print(f"y = {y} at t = {float(t):.3f} s")

Run it from the root of the repository.

uv run python examples/quickstart.py

It prints one line like this.

y = tensor([3., 5., 7.]) at t = 0.017 s

y is \(2x + 1\) for x = [1, 2, 3]. t is the time at which shift sent y, in seconds of the run's timer. It changes from run to run.

A node¶

class Affine(Node):
    """Sends ``a * x + b``, where ``x`` is the latest input."""

    def __init__(self, a: float, b: float) -> None:
        super().__init__()
        self.a, self.b = a, b  # settings, not state
        self.prev = Port(recv=VECTOR)  # receives from the previous node
        self.next = Port(send=VECTOR)  # sends to the next node

    @override
    def on_update(self, dt: float) -> None:
        super().on_update(dt)
        self.next.send(self.a * self.prev.recv() + self.b)  # recv() never waits

A node is a class that inherits from mujou.Node. As in any subclass, __init__ calls super().__init__() first, and each hook calls the same hook of super() first
__init__ declares the node's ports as attributes. Port(recv=VECTOR) receives float32 vectors of 3, and Port(send=VECTOR) sends them. By convention, the port on the side of the previous node is prev and the one on the side of the next node is next
a and b are settings: numbers set from the arguments, not state
on_update(dt) is one step of the node's work. Mujou calls it again and again, about every 10 ms. recv() returns the latest value the peer sent, without waiting, and send() sends a copy to every peer

A network¶

class Pipeline(Network):
    """``x -> scale -> shift ->``: y = 2 x + 1."""

    def __init__(self) -> None:
        super().__init__()
        self.scale = Affine(2.0, 0.0)
        self.shift = Affine(1.0, 1.0)
        self.scale.next.connect(self.shift.prev)
        self.prev, self.next = self.scale.prev, self.shift.next  # exposed as its own

A network holds nodes as attributes, connects their ports with connect, and exposes the ports that the outside uses as attributes of its own (self.prev, self.next). connect checks that the spec each port sends matches the spec the other receives, and raises TypeError where it does not.

The run¶

net = Pipeline()
x = torch.tensor([1.0, 2.0, 3.0])
# The nodes run on the run's worker threads until the end of the `with` block.
with mujou.run(net) as run:
    run.write({net.prev: x})
    # x takes a few updates to pass through the nodes: wait until y arrives, with
    # the time it was sent (TimeoutError after 5 s of real time).
    y, t = run.wait_for(net.next, lambda y: torch.equal(y, 2 * x + 1), timeout=5.0)
    assert t is not None  # None only for the zeros before the first send

print(f"y = {y} at t = {float(t):.3f} s")

mujou.run(net) checks the network, starts every node, and returns a Run once they are running. Each node repeats on_update on the run's worker threads; the host's thread goes on
run.write writes x to the unconnected port net.prev. The value stays there until the next write
recv() never waits, so shift keeps sending values computed from older inputs (zeros at first) until x has passed through scale. run.wait_for returns the first output that meets the condition, with its time, and raises TimeoutError after 5 seconds of real time
Leaving the with block shuts the run down, and an exception of a node reaches the host there at the latest

Branching and merging¶

A Port that only sends takes any number of peers: connect it to several ports, and each gets every value it sends (fan-out). A MultiPort merges: its recv() returns one value per channel, as a dict[Channel, value] in the order of connect, and the node decides how to combine them.

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

Diamond computes \(y = 2(x + 1) + (x + 1)^2\): shift sends to both branches, and sum adds what the two send. A Port that receives takes one peer only, because its recv() returns one value; receive from several with a MultiPort. Run it with uv run python examples/branching.py. The same program also sends a different value to each peer, which Ports and Channels explains.

Next¶

Write a node whose state changes over time: 2. Write a Node
The host's other ways to write and read: Host I/O