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)[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.

Returns:

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

Return type:

Sequence[HiddenParamOpRelation]