GraphDef

class brainstate.graph.GraphDef(root, node_specs)

The static structure of an object graph.

A GraphDef is a root Edge plus a flat tuple of graph-node specs ordered ascending by index. It is hashable (the hash is computed once and cached) and registered as a JAX static pytree, so it can serve as a transform cache key (e.g. for jit).

Parameters:

• root (Edge) – The edge for the root object (a NodeEdge, PytreeEdge, etc.).

• node_specs (tuple) – Graph-node specs, ascending by index.

Examples

>>> import brainstate
>>> model = brainstate.nn.Linear(2, 3)
>>> graphdef, _ = brainstate.graph.flatten(model)
>>> graphdef.type is brainstate.nn.Linear
True

property index: int

Global index of the root node, or 0 when the root is a pytree.

property type: type

Type of the root container (the root node type, or PytreeType).