def run_inference_and_populate_raw_store(
    model_context: StandardModelContext,
    multi_autoencoder_context: MultiAutoencoderContext | None,
    batched_ds_computation_params: list[DerivedScalarComputationParams],
) -> tuple[
    list[RawActivationStore],
    list[float | None],
    list[float | None],
    list[float | None],
    float,
]:
    """
    This populates a dict of ActivationsAndMetadata objects for each batch element, and returns
    inference-related stats.

     - batched_requested_activations_by_location_type_and_pass_type: stored activations
     - batched_loss_floats: loss values for each batch element, if a loss function was specified
     - batched_activation_value_for_backward_pass_floats: value of activation for which a backward pass was
    computed for each batch element, if an activation index for backward pass was specified
     - batched_memory_used_before: amount of memory allocated to the GPU before the forward pass for each batch element
     - inference_time: time used for the forward pass, in seconds
    """
    for params in batched_ds_computation_params:
        trace_config = params.trace_config
        if trace_config is not None:
            if trace_config.node_type.is_autoencoder_latent:
                assert multi_autoencoder_context is not None
                assert (
                    multi_autoencoder_context.get_autoencoder_context(trace_config.node_type)
                    is not None
                ), f"Autoencoder context not found for {trace_config.node_type}"

    transformer = model_context.get_or_create_model()
    batched_input_token_ints = [params.input_token_ints for params in batched_ds_computation_params]
    tokens_tensor, pad_tensor = prep_input_and_right_pad_for_forward_pass(
        batched_input_token_ints, transformer.device
    )

    (
        batched_upstream_activation_index_for_backward_pass,
        batched_downstream_activation_index_for_backward_pass,
        batched_reconstitute_activation_fn,
    ) = replace_activation_index_using_reconstituter(
        model_context=model_context,
        multi_autoencoder_context=multi_autoencoder_context,
        batched_ds_computation_params=batched_ds_computation_params,
    )

    batched_activation_index_for_backward_pass_by_name = [
        {
            "upstream": upstream_activation_index_for_backward_pass,
            "downstream": downstream_activation_index_for_backward_pass,
        }
        for (
            upstream_activation_index_for_backward_pass,
            downstream_activation_index_for_backward_pass,
        ) in zip(
            batched_upstream_activation_index_for_backward_pass,
            batched_downstream_activation_index_for_backward_pass,
        )
    ]
    (
        transformer_graph,  # Stores the hooks
        batched_requested_activations_by_location_type_and_pass_type,  # Stores the activations from the hooks after the forward pass
        batched_requested_attached_activations_for_backward_pass_by_name,
    ) = get_transformer_graph_hooks_and_activation_caches(
        multi_autoencoder_context=multi_autoencoder_context,
        batched_ds_computation_params=batched_ds_computation_params,
        batched_activation_index_for_backward_pass_by_name=batched_activation_index_for_backward_pass_by_name,
    )

    assert len(batched_requested_activations_by_location_type_and_pass_type) == len(
        batched_ds_computation_params
    )
    for (
        requested_activations_by_location_type_and_pass_type,
        ds_computation_params,
    ) in zip(
        batched_requested_activations_by_location_type_and_pass_type,
        batched_ds_computation_params,
    ):
        pass_types = [
            activation_location_type_and_pass_type.pass_type
            for activation_location_type_and_pass_type in requested_activations_by_location_type_and_pass_type.keys()
        ]

        if any(pass_type == PassType.BACKWARD for pass_type in pass_types):
            assert (
                ds_computation_params.loss_fn_for_backward_pass is not None
                or ds_computation_params.trace_config is not None
            ), "loss_fn_for_backward_pass or trace_config must be defined if gradients are required"

    batched_device_for_raw_activations = [
        params.device_for_raw_activations for params in batched_ds_computation_params
    ]

    t0 = time.time()
    cuda_available = torch.cuda.is_available()
    if cuda_available and any(
        device.type == "cuda" for device in batched_device_for_raw_activations
    ):
        torch.cuda.empty_cache()
    batched_memory_used_before: list[float | None] = [
        torch.cuda.memory_allocated(device) if device.type == "cuda" and cuda_available else None
        for device in batched_device_for_raw_activations
    ]

    logits, _ = transformer.forward(
        tokens_tensor, pad=pad_tensor, hooks=transformer_graph.as_transformer_hooks()
    )
    batched_loss: list[torch.Tensor | None] = []
    batched_activation_value_for_backward_pass: list[torch.Tensor | None] = []
    for batch_index, (
        ds_computation_params,
        requested_attached_activation_for_backward_pass_by_name,
        reconstitute_activation_fn,
        activation_index_for_backward_pass_by_name,
    ) in enumerate(
        zip(
            batched_ds_computation_params,
            batched_requested_attached_activations_for_backward_pass_by_name,
            batched_reconstitute_activation_fn,
            batched_activation_index_for_backward_pass_by_name,
        )
    ):
        loss_fn_for_backward_pass = ds_computation_params.loss_fn_for_backward_pass
        if loss_fn_for_backward_pass is not None:
            loss = loss_fn_for_backward_pass(logits[batch_index].unsqueeze(0))
        else:
            loss = None

        if activation_index_for_backward_pass_by_name["upstream"] is not None:
            assert requested_attached_activation_for_backward_pass_by_name["upstream"] is not None
            activation_value_for_backward_pass = reconstitute_activation_fn(
                requested_attached_activation_for_backward_pass_by_name["upstream"],
                requested_attached_activation_for_backward_pass_by_name["downstream"],
            )
        else:
            activation_value_for_backward_pass = None

        batched_loss.append(loss)
        batched_activation_value_for_backward_pass.append(activation_value_for_backward_pass)

    populated_losses: list[torch.Tensor] = []
    for loss, value in zip(batched_loss, batched_activation_value_for_backward_pass):
        # backward pass is computed from value if it is not None, otherwise from loss
        if value is not None:
            populated_losses.append(value)
        elif loss is not None:
            populated_losses.append(loss)

    if len(populated_losses):
        assert all(isinstance(loss, torch.Tensor) for loss in populated_losses)
        loss_sum = sum(populated_losses)
        assert isinstance(loss_sum, torch.Tensor)
        loss_sum.backward()

    inference_time = time.time() - t0

    batched_loss_floats: list[float | None] = [
        loss.item() if loss is not None and not torch.isnan(loss) else None for loss in batched_loss
    ]

    batched_activation_value_for_backward_pass_floats: list[float | None] = [
        activation.item() if activation is not None else None
        for activation in batched_activation_value_for_backward_pass
    ]

    batched_raw_activation_store: list[RawActivationStore] = []
    for (requested_activations_by_location_type_and_pass_type,) in zip(
        batched_requested_activations_by_location_type_and_pass_type,
    ):
        raw_activation_store = RawActivationStore.from_nested_dict_of_activations(
            requested_activations_by_location_type_and_pass_type
        )
        batched_raw_activation_store.append(raw_activation_store)

    assert (
        len(batched_raw_activation_store)
        == len(batched_loss_floats)
        == len(batched_activation_value_for_backward_pass_floats)
        == len(batched_memory_used_before)
        == len(batched_ds_computation_params)
    )  # returns one batch element per input param setting

    return (
        batched_raw_activation_store,
        batched_loss_floats,
        batched_activation_value_for_backward_pass_floats,
        batched_memory_used_before,
        inference_time,
    )

    This populates a dict of ActivationsAndMetadata objects for each batch element, and returns
    inference-related stats.

     - batched_requested_activations_by_location_type_and_pass_type: stored activations
     - batched_loss_floats: loss values for each batch element, if a loss function was specified
     - batched_activation_value_for_backward_pass_floats: value of activation for which a backward pass was
    computed for each batch element, if an activation index for backward pass was specified
     - batched_memory_used_before: amount of memory allocated to the GPU before the forward pass for each batch element
     - inference_time: time used for the forward pass, in seconds
    
neuron_explainer/activation_server/derived_scalar_computation.py