ModuleInfo

class braintrace.ModuleInfo

The model information for the ETrace compiler.

Bundles the abstract representation of a model and all the lookup tables the compiler needs. It groups information into five categories: the stateful model, the jaxpr, the states, the hidden states, and the parameter weights.

Variables:

• stateful_model (brainstate.transform.StatefulFunction) – The stateful function that compiles the model into an abstract jaxpr representation.

• closed_jaxpr (ClosedJaxpr) – The closed-jaxpr representation of the model.

• retrieved_model_states (brainstate.util.FlattedDict) – The model states retrieved from model.states(), with well-defined paths and structures.

• compiled_model_states (sequence of brainstate.State) – The model states compiled from the stateful model; accurate and consistent with the model jaxpr but lacking path information.

• state_id_to_path (dict) – Mapping from each state id to its state path.

• state_tree_invars (PyTree of Var) – The input jaxpr variables of the states, as a pytree.

• state_tree_outvars (PyTree of Var) – The output jaxpr variables of the states, as a pytree.

• hidden_path_to_invar (dict) – Mapping from each hidden path to its input variable.

• hidden_path_to_outvar (dict) – Mapping from each hidden path to its output variable.

• invar_to_hidden_path (dict) – Mapping from each input variable to its hidden path.

• outvar_to_hidden_path (dict) – Mapping from each output variable to its hidden path.

• hidden_outvar_to_invar (dict) – Mapping from each output variable to its input variable.

• weight_invars (list of Var) – The weight input variables.

• weight_path_to_invars (dict) – Mapping from each weight path to its input variables.

• invar_to_weight_path (dict) – Mapping from each input variable to its weight path.

• num_var_out (int) – Number of original output variables.

• num_var_state (int) – Number of state-variable outputs.

• control_flow (ControlFlowPolicy) –

  The control-flow policy the canonicalizer ran with. Downstream passes (structured scan descent, hidden-group discovery, relation discovery, hidden perturbation) consult this same policy so opaque control-flow handling is consistent across the whole compilation.

  Structured descent of ETP-relevant scans above the unroll limit is implemented by the internal braintrace._compiler.scan_descent module.

See also

extract_module_info

Build a ModuleInfo from a model.

Examples

>>> import brainstate
>>> import braintrace
>>> gru = braintrace.nn.GRUCell(3, 4)
>>> _ = brainstate.nn.init_all_states(gru)
>>> inputs = brainstate.random.randn(3)
>>> module_info = braintrace.extract_module_info(gru, inputs)
>>> isinstance(module_info, braintrace.ModuleInfo)
True

static __new__(_cls, stateful_model, closed_jaxpr, retrieved_model_states, compiled_model_states, state_id_to_path, state_tree_invars, state_tree_outvars, hidden_path_to_invar, hidden_path_to_outvar, invar_to_hidden_path, outvar_to_hidden_path, hidden_outvar_to_invar, weight_invars, weight_path_to_invars, invar_to_weight_path, num_var_out, num_var_state, control_flow=ControlFlowPolicy(cond='convert', scan_unroll_limit=16, while_hidden='opaque-fwd', etp_in_control_flow='error', scan_descent='auto', fixpoint_iteration_limit=64))

Create new instance of ModuleInfo(stateful_model, closed_jaxpr, retrieved_model_states, compiled_model_states, state_id_to_path, state_tree_invars, state_tree_outvars, hidden_path_to_invar, hidden_path_to_outvar, invar_to_hidden_path, outvar_to_hidden_path, hidden_outvar_to_invar, weight_invars, weight_path_to_invars, invar_to_weight_path, num_var_out, num_var_state, control_flow)

add_jaxpr_outs(jax_vars)

Add extra jaxpr outputs to the model jaxpr.

Returns a new ModuleInfo whose jaxpr additionally outputs the given variables, so the compiler can recover the intermediate values it needs.

Parameters:

jax_vars (sequence of Var) – The extra jaxpr variables to append to the jaxpr outputs.

Returns:

ModuleInfo – A new ModuleInfo with the extended jaxpr.

closed_jaxpr

Alias for field number 1

compiled_model_states

Alias for field number 3

control_flow

Alias for field number 17

dict()

Return this module info’s named fields as a plain dictionary.

Returns:

dict – An ordered mapping from field name to value, as produced by the underlying typing.NamedTuple.

hidden_outvar_to_invar

Alias for field number 11

hidden_path_to_invar

Alias for field number 7

hidden_path_to_outvar

Alias for field number 8

invar_to_hidden_path

Alias for field number 9

invar_to_weight_path

Alias for field number 14

property jaxpr

The jaxpr of the model.

Returns:

Jaxpr – The jaxpr extracted from closed_jaxpr.

jaxpr_call(*args, old_state_vals=None)

Evaluate the model on the given inputs using the compiled jaxpr.

Parameters:

• *args (Inputs) – The inputs of the model.

• old_state_vals (sequence of jax.Array or None, optional) – The old state values. When None, the current values of the compiled model states are used. Default None.

Returns:

• out (Outputs) – The output of the model.

• etrace_vals (ETraceVals) – The values for the eligibility-trace (hidden) states.

• oth_state_vals (StateVals) – The other state values.

• temps (TempData) – The temporary intermediate values.

num_var_out

Alias for field number 15

num_var_state

Alias for field number 16

outvar_to_hidden_path

Alias for field number 10

retrieved_model_states

Alias for field number 2

split_state_outvars()

Split the state outvars into weight, hidden, and other states.

Returns:

• weight_jaxvar_tree (PyTree of Var) – The weight tree of jaxpr variables.

• hidden_jaxvar (PyTree of Var) – The hidden tree of jaxpr variables.

• other_state_jaxvar_tree (PyTree of Var) – The other-state tree of jaxpr variables.

state_id_to_path

Alias for field number 4

state_tree_invars

Alias for field number 5

state_tree_outvars

Alias for field number 6

stateful_model

Alias for field number 0

weight_invars

Alias for field number 12

weight_path_to_invars

Alias for field number 13

ModuleInfo.__init__()