extract_module_info

extract_module_info#

class braintrace.extract_module_info(model, *model_args, control_flow=None, **model_kwargs)[source]#

Extract the model information for the ETrace compiler.

Parameters:

  • model (Module) – The model from which to extract the information.

  • *model_args – The positional arguments of the model.

  • control_flow (ControlFlowPolicy | None) – Policy governing control-flow canonicalization (cond if-conversion and inner-scan unrolling; see canonicalize_control_flow()) and downstream handling of un-flattened control flow. None (default) uses the default policy, which converts every ETP-relevant cond, unrolls every ETP-relevant scan of static length at most 16, keeps weight-free while loops that touch hidden state as opaque forward nodes (while_hidden='opaque-fwd'), and raises on ETP primitives left inside a control-flow body (etp_in_control_flow='error'). The policy is stored on the returned ModuleInfo (minfo.control_flow) so later compiler passes apply the same rules.

  • **model_kwargs – The keyword arguments of the model.

Returns:

The model information.

Return type:

ModuleInfo

See also

ModuleInfo

The returned data structure.

Notes

Prefer positional arguments. **model_kwargs is accepted here for tracing, but ModuleInfo.jaxpr_call and the downstream compile_etrace_graph pipeline rebuild inputs from positional arguments only — bind static keyword arguments with functools.partial before compiling.

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)
>>> module_info.num_var_out
1