find_hidden_param_op_relations_from_minfo

find_hidden_param_op_relations_from_minfo#

class braintrace.find_hidden_param_op_relations_from_minfo(minfo, hid_path_to_group, descended_scan_eqn_ids=frozenset({}))[source]#

Find ETP relations from a ModuleInfo.

Builds a mapping from all brainstate.ParamState invars so that plain ParamState weights used with ETP primitives are recognised.

Parameters:

minfo (ModuleInfo) – The model information.
hid_path_to_group (Dict[Tuple[str, ...], HiddenGroup]) – Mapping from each hidden-state path to its HiddenGroup.
descended_scan_eqn_ids (FrozenSet[int]) – id() values of scan equations rewritten by structured scan descent (Phase 4); those equations are skipped by the relation walker (their body relations are registered by the scan-descent pass).

Returns:

The discovered ETP-primitive-to-hidden-state relations.

Return type:

Sequence[HiddenParamOpRelation]