braintrace.find_hidden_param_op_relations_from_module

braintrace.find_hidden_param_op_relations_from_module#

braintrace.find_hidden_param_op_relations_from_module(model, *model_args, **model_kwargs)#

Find ETP relations from a model.

Parameters:

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

  • *model_args – The positional arguments of the model.

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

Returns:

sequence of HiddenParamOpRelation – The discovered ETP-primitive-to-hidden-state relations.

See also

find_hidden_param_op_relations_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)
>>> relations = braintrace.find_hidden_param_op_relations_from_module(gru, inputs)
>>> len(relations)
2