async def test_activation_grabbing(
    standard_model_context: StandardModelContext,
    standard_autoencoder_context: AutoencoderContext,
) -> None:
    """Compute act-times-grad of all upstream nodes for a given node"""
    interactive_model = InteractiveModel.from_standard_model_context_and_autoencoder_context(
        standard_model_context,
        standard_autoencoder_context,
    )

    # specify from what points the backward pass will be performed; either from the loss or
    # from an activation, running each backward pass one after another
    layer_index_for_grad = 6
    inference_kwargs_by_backward_pass_setting: dict[str, dict[str, Any]] = {
        "loss": {
            "loss_fn_config": LossFnConfig(
                name=LossFnName.LOGIT_DIFF,
                target_tokens=["!"],
                distractor_tokens=["."],
            ),
            "trace_config": None,
        },
        "activation": {
            "loss_fn_config": None,
            "trace_config": MirroredTraceConfig.from_trace_config(
                TraceConfig.from_activation_index(
                    ActivationIndex(
                        layer_index=layer_index_for_grad,
                        activation_location_type=ActivationLocationType.MLP_POST_ACT,
                        tensor_indices=(0, 1),
                        pass_type=PassType.FORWARD,
                    ),
                    detach_layer_norm_scale=DETACH_LAYER_NORM_SCALE_FOR_TEST,
                )
            ),
        },
        "autoencoder_activation": {
            "loss_fn_config": None,
            "trace_config": MirroredTraceConfig.from_trace_config(
                TraceConfig.from_activation_index(
                    ActivationIndex(
                        layer_index=layer_index_for_grad,
                        activation_location_type=ActivationLocationType.ONLINE_AUTOENCODER_LATENT,
                        tensor_indices=(0, 1),
                        pass_type=PassType.FORWARD,
                    ),
                    detach_layer_norm_scale=DETACH_LAYER_NORM_SCALE_FOR_TEST,
                )
            ),
        },
    }

    # For each node, compute the gradient from this node to all upstream nodes, and select the
    # top-k upstream nodes with the largest activation-times-gradient.
    grads_by_backward_pass_setting = {}  # more precisely "act-times-grad"
    for backward_pass_setting in [
        "loss",
        "activation",
        "autoencoder_activation",
    ]:
        inference_kwargs: dict[str, Any] = inference_kwargs_by_backward_pass_setting[
            backward_pass_setting
        ]
        grad_requiring_dst_request_spec = MultipleTopKDerivedScalarsRequestSpec(
            dst_list_by_group_id=make_dst_list_by_group_id(
                group_ids=[
                    GroupId.ACT_TIMES_GRAD,
                ],
                component_types=[
                    "mlp",
                    "autoencoder",
                    "unflattened_attn",
                ],
            ),
            # dsts for each group ID are assumed to have defined node_type,
            # all node_types assumed to be distinct within a group_id, and all group_ids to
            # contain the same set of node_types.
            token_index=None,
            top_and_bottom_k=10,
            pass_type=PassType.FORWARD,
        )
        multiple_top_k_layers_request_spec = MultipleTopKDerivedScalarsRequestSpec(
            dst_list_by_group_id={
                GroupId.SINGLETON: [
                    DerivedScalarType.RESID_POST_ATTN_PROJ_TO_FINAL_RESIDUAL_GRAD,
                ]
            },
            # dsts for each group ID are assumed to have defined node_type,
            # all node_types assumed to be distinct within a group_id, and all group_ids to
            # contain the same set of node_types.
            token_index=None,
            top_and_bottom_k=100,
            pass_type=PassType.FORWARD,
        )
        print(f"Testing backward pass setting: {backward_pass_setting}")
        batched_request = BatchedRequest(
            inference_sub_requests=[
                InferenceSubRequest(
                    inference_request_spec=InferenceRequestSpec(
                        prompt="Hello world",
                        **inference_kwargs,
                    ),
                    processing_request_spec_by_name={
                        "multi_top_k": grad_requiring_dst_request_spec,
                        "layers": multiple_top_k_layers_request_spec,
                    },
                ),
                InferenceSubRequest(
                    inference_request_spec=InferenceRequestSpec(
                        prompt="Goodbye",
                        **inference_kwargs,
                    ),
                    processing_request_spec_by_name={
                        "multi_top_k": grad_requiring_dst_request_spec,
                        "layers": multiple_top_k_layers_request_spec,
                    },
                ),
            ],
        )
        assert (
            batched_request.inference_sub_requests[0].inference_request_spec.trace_config
            is not None
            or batched_request.inference_sub_requests[0].inference_request_spec.loss_fn_config
            is not None
        )
        batched_response = await interactive_model.handle_batched_request(batched_request)
        # We add two requests, so we should get two responses. For now, only check
        # the validity of the first response though.
        assert len(batched_response.inference_sub_responses) == 2
        response = batched_response.inference_sub_responses[0]
        multi_top_k = response.processing_response_data_by_name["multi_top_k"]
        assert isinstance(multi_top_k, MultipleTopKDerivedScalarsResponseData)
        grads_by_backward_pass_setting[backward_pass_setting] = multi_top_k.activations_by_group_id[
            GroupId.ACT_TIMES_GRAD
        ]

    # Check that the values are different for different backward pass settings
    assert not np.array_equal(
        grads_by_backward_pass_setting["loss"],
        grads_by_backward_pass_setting["activation"],
    )
    assert not np.array_equal(
        grads_by_backward_pass_setting["activation"],
        grads_by_backward_pass_setting["autoencoder_activation"],
    )

    # Checking the extremal values of each backward pass setting
    # Note: will need to change if anything about the requests is changed
    top_value_by_backward_pass_setting = {
        "loss": 162.00540161132812,
        "activation": 0.19626690447330475,
        "autoencoder_activation": 0.44310837984085083,
    }
    for backward_pass_setting in [
        "loss",
        "activation",
        "autoencoder_activation",
    ]:
        grads = grads_by_backward_pass_setting[backward_pass_setting]
        assert math.isclose(
            grads[0], top_value_by_backward_pass_setting[backward_pass_setting], rel_tol=1e-4
        ), f"{grads[0]} != {top_value_by_backward_pass_setting[backward_pass_setting]}"
Compute act-times-grad of all upstream nodes for a given node
neuron_explainer/tests/test_interactive_model.py