def truncate_to_expected_shape_tensor_calculate_derived_scalar_fn(
    raw_activation_data_tuple: tuple[torch.Tensor, ...],
    layer_index: LayerIndex,
    pass_type: PassType,
    expected_shape: list[int | None],
) -> torch.Tensor:
    """This either:
    converts a length 1 tuple of tensors into
    a single tensor; pass_type is asserted to be PassType.FORWARD
    or
    converts a length 2 tuple of tensors, one for the forward pass and one for the backward pass,
    into the appropriate one of those two objects, depending on the pass_type argument."""
    if len(raw_activation_data_tuple) == 1:
        # in this case, only the activations at the relevant ActivationLocationType have been loaded from disk
        assert pass_type == PassType.FORWARD
        raw_activation_data = raw_activation_data_tuple[0]
        raw_activation_data = truncate_to_expected_shape(
            raw_activation_data,
            expected_shape,
        )
        return raw_activation_data
    elif len(raw_activation_data_tuple) == 2:
        # in this case, both the activations and gradients at the relevant ActivationLocationType have been loaded from disk
        raw_activation_data, raw_gradient_data = raw_activation_data_tuple
        if pass_type == PassType.FORWARD:
            raw_activation_data = truncate_to_expected_shape(
                raw_activation_data,
                expected_shape,
            )
            return raw_activation_data
        elif pass_type == PassType.BACKWARD:
            raw_gradient_data = truncate_to_expected_shape(
                raw_gradient_data,
                expected_shape,
            )
            return raw_gradient_data
        else:
            raise ValueError(f"Unknown {pass_type=}")
    else:
        raise ValueError(f"Unknown {raw_activation_data_tuple=}")
This either:
    converts a length 1 tuple of tensors into
    a single tensor; pass_type is asserted to be PassType.FORWARD
    or
    converts a length 2 tuple of tensors, one for the forward pass and one for the backward pass,
    into the appropriate one of those two objects, depending on the pass_type argument.
neuron_explainer/activations/derived_scalars/raw_activations.py