def make_attn_write_to_latent_per_sequence_token_batched_scalar_deriver(
    dst_config: DstConfig,
) -> ScalarDeriver:
    """Same as make_attn_write_to_latent_per_sequence_token_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_tokens, n_heads, n_latents).
    """
    autoencoder_node_type = NodeType.ATTENTION_AUTOENCODER_LATENT
    attn_weighted_sum_of_value_scalar_deriver = (
        make_scalar_deriver_factory_for_activation_location_type(
            ActivationLocationType.ATTN_WEIGHTED_SUM_OF_VALUES
        )(dst_config)
    )  # derive scalar of shape (n_tokens, n_heads, v_channels)

    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
    )

    model_context = dst_config.get_model_context()
    W_O = model_context.get_weight(
        location_type=WeightLocationType.ATTN_TO_RESIDUAL,
        layer=layer_index,
        device=model_context.device,
    )  # shape (n_heads, d_head, d_model)

    def attn_write_to_latents_tensor_calculate_derived_scalar_fn(
        attn_weighted_sum_of_value: 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_heads, d_head = attn_weighted_sum_of_value.shape
        assert autoencoder_latents.shape[0] == n_tokens
        n_tokens, n_latents = autoencoder_latents.shape
        assert W_O.shape[:2] == (n_heads, d_head)
        n_heads, d_head, d_model = W_O.shape

        # from (n_heads, d_head, d_model) to (n_heads, d_head, n_latents)
        Wo_encoder = pre_act_encoder_derivative(W_O)
        Wo_encoder = Wo_encoder.to(attn_weighted_sum_of_value.dtype)

        projection = torch.einsum("thd,hdl->thl", attn_weighted_sum_of_value, Wo_encoder)

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

    new_scalar_deriver = attn_weighted_sum_of_value_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.ATTN_WRITE_TO_LATENT_PER_SEQUENCE_TOKEN_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_heads, n_latents)
Same as make_attn_write_to_latent_per_sequence_token_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_tokens, n_heads, n_latents).
    
neuron_explainer/activations/derived_scalars/attention.py