def make_flattened_attn_write_to_latent_summed_over_heads_batched_scalar_deriver(
    dst_config: DstConfig,
) -> ScalarDeriver:
    """Same as make_flattened_attn_write_to_latent_summed_over_heads_scalar_deriver, but over
    all latents at the same time.
    Assumes that the autoencoder has a single encoder layer and a ReLU, to compute the
    gradient (attribution) without backprop.
    Output has shape (n_token_pairs, n_latents).
    """
    autoencoder_node_type = NodeType.ATTENTION_AUTOENCODER_LATENT
    attn_write_sum_heads_scalar_deriver = make_attn_write_sum_heads_scalar_deriver(
        dst_config
    )  # derive scalar of shape (n_tokens, n_attended_to_tokens, d_model)
    autoencoder_scalar_deriver = make_autoencoder_latent_scalar_deriver_factory(
        autoencoder_node_type
    )(
        dst_config
    )  # derive scalar of shape (n_tokens, n_latents)

    assert dst_config.layer_indices is not None
    assert len(dst_config.layer_indices) == 1
    layer_index = dst_config.layer_indices[0]

    autoencoder_context = dst_config.get_autoencoder_context(autoencoder_node_type)
    assert autoencoder_context is not None
    pre_act_encoder_derivative = make_autoencoder_pre_act_encoder_derivative(
        autoencoder_context, layer_index
    )
    activation_fn_derivative = make_autoencoder_activation_fn_derivative(
        autoencoder_context, layer_index
    )

    def attn_write_to_latents_tensor_calculate_derived_scalar_fn(
        attn_write_sum_heads: torch.Tensor,
        autoencoder_latents: torch.Tensor,
        layer_index: LayerIndex,
        pass_type: PassType,
    ) -> torch.Tensor:
        assert layer_index == layer_index
        assert pass_type == PassType.FORWARD, "write to latents only defined for forward pass"
        n_tokens, n_tokens_attended_to, d_model = attn_write_sum_heads.shape
        assert autoencoder_latents.shape[0] == n_tokens
        n_tokens, n_latents = autoencoder_latents.shape
        # from (n_tokens, n_tokens_attended_to, d_model) to (n_tokens, n_tokens_attended_to, n_latents)
        projection = pre_act_encoder_derivative(attn_write_sum_heads)

        d_latent_d_pre_act = activation_fn_derivative(autoencoder_latents)
        direct_write_to_latents = torch.einsum("tul,tl->tul", projection, d_latent_d_pre_act)

        flattened_direct_write_to_latents = (
            _convert_rectangular_activations_to_flattened_lower_triangle(direct_write_to_latents)
        )
        return flattened_direct_write_to_latents

    new_scalar_deriver = (
        attn_write_sum_heads_scalar_deriver.apply_layerwise_transform_fn_to_output_and_other_tensor(
            layerwise_transform_fn=attn_write_to_latents_tensor_calculate_derived_scalar_fn,
            pass_type_to_transform=PassType.FORWARD,
            output_dst=DerivedScalarType.FLATTENED_ATTN_WRITE_TO_LATENT_SUMMED_OVER_HEADS_BATCHED,
            other_scalar_source=DerivedScalarSource(
                scalar_deriver=autoencoder_scalar_deriver,
                pass_type=PassType.FORWARD,
                layer_indexer=ConstantLayerIndexer(layer_index),
            ),
        )
    )
    return new_scalar_deriver  # derive scalar of shape (n_token_pairs, n_latents)
Same as make_flattened_attn_write_to_latent_summed_over_heads_scalar_deriver, but over
    all latents at the same time.
    Assumes that the autoencoder has a single encoder layer and a ReLU, to compute the
    gradient (attribution) without backprop.
    Output has shape (n_token_pairs, n_latents).
    
neuron_explainer/activations/derived_scalars/attention.py