Node

class brainstate.graph.Node(*args, **kwargs)

Base class for all mutable graph nodes in the BrainState framework.

Subclasses are automatically registered with the graph engine, so their attributes participate in treefy_split() / treefy_merge(), state/node introspection, and the functional transforms built on top of them. Attributes are flattened in alphabetical order; names listed in graph_invisible_attrs are excluded.

Examples

>>> import brainstate
>>> class MyModel(brainstate.graph.Node):
...     def __init__(self):
...         self.w = brainstate.ParamState(brainstate.random.randn(3))
>>> model = MyModel()
>>> graphdef, params = brainstate.graph.treefy_split(model)
>>> rebuilt = brainstate.graph.treefy_merge(graphdef, params)
>>> isinstance(rebuilt, MyModel)
True

check_valid_context(error_msg)

Raise TraceContextError if the node’s trace context is invalid.

A graph node carries no trace state of its own; the only objects that track a JAX trace level are the States it holds. The node is therefore valid exactly when every reachable State is valid. If any contained State was created at an incompatible trace level, error_msg() is raised.

Parameters:

error_msg (Callable[[], str]) – Zero-argument callable returning the error message; evaluated lazily only when a violation is detected.

Raises:

TraceContextError – If any State reachable from this node has an invalid trace.

Return type:

None