register_graph_node_type

class brainstate.graph.register_graph_node_type(type, flatten, set_key, pop_key, create_empty, clear)

Register a custom mutable graph-node type with the graph engine.

Parameters:

• type (type) – The Python type to register.

• flatten (Callable[[TypeVar(N)], tuple[Sequence[tuple[Key, TypeVar(Leaf)]], TypeVar(AuxData)]]) – node -> (items, metadata) where items is an ordered sequence of (key, value) pairs and metadata is hashable auxiliary data.

• set_key (Callable[[TypeVar(N), Key, TypeVar(Leaf)], None]) – (node, key, value) -> None setting one child in place.

• pop_key (Callable[[TypeVar(N), Key], TypeVar(Leaf)]) – (node, key) -> value removing and returning one child.

• create_empty (Callable[[TypeVar(AuxData)], TypeVar(N)]) – metadata -> node building an empty shell of the type.

• clear (Callable[[TypeVar(N)], None]) – node -> None removing all children in place.

Return type:

None

Examples

>>> import brainstate
>>> class MyList:
...     def __init__(self):
...         self.items = {}
>>> brainstate.graph.register_graph_node_type(
...     MyList,
...     flatten=lambda n: (sorted(n.items.items()), (MyList,)),
...     set_key=lambda n, k, v: n.items.__setitem__(k, v),
...     pop_key=lambda n, k: n.items.pop(k),
...     create_empty=lambda meta: MyList(),
...     clear=lambda n: n.items.clear(),
... )