def make_reconstituted_gradient_fn(
    transformer: Transformer,
    autoencoder_context: AutoencoderContext | None,
    trace_config: TraceConfig,
) -> Callable[[torch.Tensor, LayerIndex, PassType], torch.Tensor]:
    """
    Define a function which takes in a 2-d residual stream tensor (along with layer index and pass
    type) and returns the gradient with respect to trace_config at those residual stream values.
    This function is asserted to be applied at the layer index and pass type specified by
    trace_config.
    """

    assert trace_config.layer_index is not None

    reconstituter = ActivationReconstituter.from_trace_config(
        transformer=transformer,
        autoencoder_context=autoencoder_context,
        trace_config=trace_config,
    )

    return reconstituter.make_reconstitute_gradient_fn_for_trace_config(
        trace_config=trace_config,
    )

    Define a function which takes in a 2-d residual stream tensor (along with layer index and pass
    type) and returns the gradient with respect to trace_config at those residual stream values.
    This function is asserted to be applied at the layer index and pass type specified by
    trace_config.
    
neuron_explainer/activations/derived_scalars/reconstituter_class.py