def convert_scalar_deriver_to_write_norm(
    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 norm for each neuron at each token.
    """

    check_write_tensor_device_matches(
        scalar_deriver,
        write_tensor_by_layer_index,
    )

    write_norm_by_layer_index = {
        layer_index_: write_tensor_by_layer_index[layer_index_].norm(dim=-1)  # type: ignore
        for layer_index_ in write_tensor_by_layer_index.keys()
    }

    def multiply_by_write_norm(
        activations: torch.Tensor,
        layer_index: LayerIndex,
        pass_type: PassType,
    ) -> torch.Tensor:
        assert pass_type == PassType.FORWARD, "Backward pass not defined for write norm"
        assert (
            layer_index in write_tensor_by_layer_index
        ), f"{layer_index=} not in {write_tensor_by_layer_index.keys()=} for {output_dst=}"
        return activations * write_norm_by_layer_index[layer_index]

    return scalar_deriver.apply_layerwise_transform_fn_to_output(
        multiply_by_write_norm,
        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 norm for each neuron at each token.
    
neuron_explainer/activations/derived_scalars/raw_activations.py