braintrace.add_hidden_perturbation_from_minfo

braintrace.add_hidden_perturbation_from_minfo#

braintrace.add_hidden_perturbation_from_minfo(minfo, descended_scan_eqn_ids=frozenset({}))#

Add hidden-state perturbations from a ModuleInfo.

Adds perturbations to the hidden states in the module jaxpr and replaces the hidden states with the perturbed states.

Parameters:
  • minfo (ModuleInfo) – The model information.

  • descended_scan_eqn_ids (frozenset of int, default frozenset()) – id() values of scan equations rewritten by structured scan descent (Phase 4); exempt from the unsupported-op checks.

Returns:

HiddenPerturbation – The hidden-perturbation information.

See also

add_hidden_perturbation_in_module

Equivalent helper starting from a model.