def get_batched_derived_scalars_for_prompt(
    model_context: StandardModelContext,
    batched_dst_and_config_list: list[
        list[tuple[DerivedScalarType, DstConfig | None]]
    ],  # None -> default config for dst
    batched_prompt: list[str],
    loss_fn_for_backward_pass: Callable[[torch.Tensor], torch.Tensor] | None = None,
    trace_config: TraceConfig | None = None,
    autoencoder_context: MultiAutoencoderContext | AutoencoderContext | None = None,
    ablation_specs: list[AblationSpec] = [],
) -> tuple[list[DerivedScalarStore], list[InferenceAndTokenData], list[RawActivationStore]]:
    """
    Lightweight function to populate a DerivedScalarStore given information specifying the prompt, loss function, and derived scalars to compute.
    """
    multi_autoencoder_context = MultiAutoencoderContext.from_context_or_multi_context(
        autoencoder_context
    )

    assert len(batched_dst_and_config_list) == len(batched_prompt)

    batched_ds_computation_params = [
        get_ds_computation_params_for_prompt(
            model_context=model_context,
            autoencoder_context=autoencoder_context,
            dst_and_config_list=dst_and_config_list,
            prompt=prompt,
            loss_fn_for_backward_pass=loss_fn_for_backward_pass,
            trace_config=trace_config,
            ablation_specs=ablation_specs,
        )
        for prompt, dst_and_config_list in zip(batched_prompt, batched_dst_and_config_list)
    ]

    batched_multi_group_ds_store_by_processing_step: list[dict[str, MultiGroupDerivedScalarStore]]
    (
        batched_multi_group_ds_store_by_processing_step,
        batched_inference_data,
        batched_raw_activation_store,
    ) = compute_derived_scalar_groups_for_input_token_ints(
        model_context=model_context,
        multi_autoencoder_context=multi_autoencoder_context,
        batched_ds_computation_params=batched_ds_computation_params,
    )
    batched_ds_store = [
        multi_group_ds_store["dummy"].to_single_ds_store()
        for multi_group_ds_store in batched_multi_group_ds_store_by_processing_step
    ]
    batched_inference_and_token_data = []
    for ds_computation_params, inference_data in zip(
        batched_ds_computation_params, batched_inference_data
    ):
        input_token_ints = ds_computation_params.input_token_ints
        input_token_strings = [
            model_context.decode_token(token_int) for token_int in input_token_ints
        ]
        batched_inference_and_token_data.append(
            InferenceAndTokenData(
                **inference_data.dict(),
                tokens_as_ints=input_token_ints,
                tokens_as_strings=input_token_strings,
            )
        )

    return batched_ds_store, batched_inference_and_token_data, batched_raw_activation_store

    Lightweight function to populate a DerivedScalarStore given information specifying the prompt, loss function, and derived scalars to compute.
    
neuron_explainer/activation_server/derived_scalar_computation.py