def convert_scalar_deriver_to_write(
    scalar_deriver: ScalarDeriver,
    write_tensor_by_layer_index: dict[LayerIndex, torch.Tensor] | dict[int, torch.Tensor],
    output_dst: DerivedScalarType,
) -> ScalarDeriver:
    """Converts a scalar deriver for a scalar type that is 1:1 with an ActivationLocationType
    to a scalar deriver for the write vector of the layer at each token."""

    check_write_tensor_device_matches(
        scalar_deriver,
        write_tensor_by_layer_index,
    )

    def multiply_by_write(
        activations: torch.Tensor,
        layer_index: LayerIndex,
        pass_type: PassType,
    ) -> torch.Tensor:
        assert pass_type == PassType.FORWARD, "Backward pass not defined for write"
        assert (
            layer_index in write_tensor_by_layer_index
        ), f"{layer_index=} not in {write_tensor_by_layer_index.keys()=}"
        return torch.einsum(
            "ta,ao->to",
            activations,
            write_tensor_by_layer_index[layer_index],  # type: ignore
        )

    return scalar_deriver.apply_layerwise_transform_fn_to_output(
        multiply_by_write,
        pass_type_to_transform=PassType.FORWARD,
        output_dst=output_dst,
    )
Converts a scalar deriver for a scalar type that is 1:1 with an ActivationLocationType
    to a scalar deriver for the write vector of the layer at each token.
neuron_explainer/activations/derived_scalars/raw_activations.py