find_hidden_groups_from_minfo

find_hidden_groups_from_minfo#

class braintrace.find_hidden_groups_from_minfo(minfo, include_recurrent_mixing=False, descended_scan_eqn_ids=frozenset({}), descended_hidden_paths=frozenset({}))[source]#

Find the hidden groups from the model information.

Parameters:

  • minfo (ModuleInfo) – The model information.

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

  • descended_scan_eqn_ids (FrozenSet[int]) – 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[Tuple[str, ...]]) – 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.