def convert_scalar_deriver_to_write_to_direction_with_write_tensor(
    scalar_deriver: ScalarDeriver,
    write_tensor_by_layer_index: dict[LayerIndex, torch.Tensor] | dict[int, torch.Tensor] | None,
    direction_scalar_source: ScalarSource,
    output_dst: DerivedScalarType,
) -> ScalarDeriver:
    """Takes as input a scalar deriver for a scalar activation fully defining a write direction
    (e.g. MLP activation or autoencoder but not post-softmax attention) and a scalar deriver for a direction
    in the residual stream basis. Multiplies each activation by its associated write vector and projects to the direction
    of interest."""
    if write_tensor_by_layer_index is not None:
        check_write_tensor_device_matches(
            scalar_deriver,
            write_tensor_by_layer_index,
        )

    def derived_scalar_precedes_direction_layer(layer_index: LayerIndex) -> bool:
        return precedes_final_layer(
            final_residual_location_within_layer=direction_scalar_source.location_within_layer,
            final_residual_layer_index=direction_scalar_source.layer_index,
            derived_scalar_location_within_layer=scalar_deriver.location_within_layer,
            derived_scalar_layer_index=layer_index,
        )

    write_to_direction_tensor_fn = make_write_to_direction_tensor_fn(
        node_type=scalar_deriver.dst.node_type,
        write_tensor_by_layer_index=write_tensor_by_layer_index,
        layer_precedes_direction_layer_fn=derived_scalar_precedes_direction_layer,
    )

    return scalar_deriver.apply_layerwise_transform_fn_to_output_and_other_tensor(
        write_to_direction_tensor_fn,
        pass_type_to_transform=PassType.FORWARD,
        output_dst=output_dst,
        other_scalar_source=direction_scalar_source,
    )
Takes as input a scalar deriver for a scalar activation fully defining a write direction
    (e.g. MLP activation or autoencoder but not post-softmax attention) and a scalar deriver for a direction
    in the residual stream basis. Multiplies each activation by its associated write vector and projects to the direction
    of interest.
neuron_explainer/activations/derived_scalars/direct_effects.py