def _attn_write_norm_tensor_calculate_derived_scalar_fn(
        raw_activation_data_tuple: tuple[torch.Tensor, ...],
        layer_index: LayerIndex,
        pass_type: PassType,
    ) -> torch.Tensor:
        """The returned processing hook, to be passed as an argument to the ScalarDeriver."""

        if len(raw_activation_data_tuple) == 2:
            attn_post_softmax_data, attn_value_data = raw_activation_data_tuple
            assert (
                pass_type == PassType.FORWARD
            ), f"Need gradient of attn_post_softmax for backward pass, got {raw_activation_data_tuple}"
        else:
            assert len(raw_activation_data_tuple) == 3
            # first position is forward pass attn_post_softmax,
            # second is backward pass attn_post_softmax,
            # third is attn_value
            if pass_type == PassType.FORWARD:
                attn_post_softmax_data, _, attn_value_data = raw_activation_data_tuple
            else:
                assert pass_type == PassType.BACKWARD
                _, attn_post_softmax_data, attn_value_data = raw_activation_data_tuple

        assert layer_index in W_O_by_layer_index

        return _compute_attn_write_norm_from_attn_and_value(
            attn_post_softmax=attn_post_softmax_data,
            attn_value=attn_value_data,
            W_O=W_O_by_layer_index[layer_index],
            pass_type=pass_type,
        )
The returned processing hook, to be passed as an argument to the ScalarDeriver.
neuron_explainer/activations/derived_scalars/attention.py