def compute_derived_scalar_groups_for_input_token_ints(
    model_context: StandardModelContext,
    multi_autoencoder_context: MultiAutoencoderContext | None,
    batched_ds_computation_params: list[DerivedScalarComputationParams],
) -> tuple[
    list[dict[str, MultiGroupDerivedScalarStore]], list[InferenceData], list[RawActivationStore]
]:
    """This function runs a batched forward pass on the given batch of input token sequences, with
    hooks added to the transformer to extract the activations needed to compute the scalars in
    multi_group_scalar_derivers for each batch element. It then returns a batch of dicts of
    DerivedScalarStores by group_id containing the relevant derived scalars for each token in
    the input, as well as a batch of InferenceData objects containing tokenized inputs and other metadata,
    and a batch of RawActivationStores, each of which was used to compute the respective
    dict of DerivedScalarStores. These RawActivationStores can be used to compute additional derived scalars
    post-hoc.
    """

    (
        batched_raw_activation_store,
        batched_loss,
        batched_activation_value_for_backward_pass,
        batched_memory_used_before,
        inference_time,
    ) = run_inference_and_populate_raw_store(
        model_context=model_context,
        multi_autoencoder_context=multi_autoencoder_context,
        batched_ds_computation_params=batched_ds_computation_params,
    )

    assert (
        len(batched_raw_activation_store)
        == len(batched_ds_computation_params)
        == len(batched_loss)
        == len(batched_activation_value_for_backward_pass)
        == len(batched_memory_used_before)
    )

    batched_multi_group_ds_store_by_processing_step: list[
        dict[str, MultiGroupDerivedScalarStore]
    ] = []
    (
        batched_multi_group_ds_store_by_processing_step,
        batched_memory_used_after,
    ) = construct_ds_stores_from_raw(
        batched_raw_activation_store,
        batched_ds_computation_params,
    )

    batched_inference_data: list[InferenceData] = []
    for (
        loss,
        activation_value_for_backward_pass,
        memory_used_before,
        memory_used_after,
    ) in zip(
        batched_loss,
        batched_activation_value_for_backward_pass,
        batched_memory_used_before,
        batched_memory_used_after,
    ):
        inference_data = InferenceData(
            inference_time=inference_time,
            loss=loss,
            activation_value_for_backward_pass=activation_value_for_backward_pass,
            memory_used_before=memory_used_before,
            memory_used_after=memory_used_after,
        )
        batched_inference_data.append(inference_data)
    return (
        batched_multi_group_ds_store_by_processing_step,
        batched_inference_data,
        batched_raw_activation_store,
    )
This function runs a batched forward pass on the given batch of input token sequences, with
    hooks added to the transformer to extract the activations needed to compute the scalars in
    multi_group_scalar_derivers for each batch element. It then returns a batch of dicts of
    DerivedScalarStores by group_id containing the relevant derived scalars for each token in
    the input, as well as a batch of InferenceData objects containing tokenized inputs and other metadata,
    and a batch of RawActivationStores, each of which was used to compute the respective
    dict of DerivedScalarStores. These RawActivationStores can be used to compute additional derived scalars
    post-hoc.
    
neuron_explainer/activation_server/derived_scalar_computation.py