braintrace.find_hidden_groups_from_minfo

braintrace.find_hidden_groups_from_minfo(minfo, include_recurrent_mixing=False, sparse_n=None, snap_max_jacobian_elements=16777216, descended_scan_eqn_ids=frozenset({}), descended_hidden_paths=frozenset({}))

Find the hidden groups from the model information.

Parameters:

• minfo (ModuleInfo) – The model information.

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

• sparse_n (int, optional) – SnAp order for recurrence_scope='sparse_n'. When given, each group carries the derived n-step neighbourhood in its snap field. Default None.

• descended_scan_eqn_ids (frozenset of int, default frozenset()) – id() values of scan equations rewritten by structured scan descent (Phase 4); those equations are skipped by the hidden-group walker.

• descended_hidden_paths (frozenset, default frozenset()) – Hidden-state paths covered by descended scan bodies; excluded from the zero-recurrence fallback grouping.

Returns:

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

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

See also

find_hidden_groups_from_module

Equivalent helper starting from a model.