braintrace.find_hidden_groups_from_module

braintrace.find_hidden_groups_from_module#

braintrace.find_hidden_groups_from_module(model, *model_args, include_recurrent_mixing=False, sparse_n=None, snap_max_jacobian_elements=16777216, **model_kwargs)#

Find hidden groups from a model.

Parameters:

  • model (brainstate.nn.Module) – The model.

  • *model_args – The positional arguments of the model.

  • include_recurrent_mixing (bool, default False) – Whether to trace recurrent ETP mixing primitives into the transition jaxpr. Keyword-only. See the internal find_hidden_groups_from_jaxpr helper for the full semantics.

  • sparse_n (int, optional) – SnAp order for recurrence_scope='sparse_n'. Keyword-only. Default None.

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

Returns:

  • hidden_groups (sequence of HiddenGroup) – The hidden groups.

  • hid_path_to_group (brainstate.util.PrettyDict) – Mapping from each hidden-state path to its HiddenGroup.

See also

find_hidden_groups_from_minfo

Equivalent helper starting from ModuleInfo.

Examples

>>> import brainstate
>>> import braintrace
>>> gru = braintrace.nn.GRUCell(3, 4)
>>> _ = brainstate.nn.init_all_states(gru)
>>> inputs = brainstate.random.randn(3)
>>> hidden_groups, hid_path_to_group = braintrace.find_hidden_groups_from_module(gru, inputs)
>>> len(hidden_groups)
1