def make_write_to_direction_tensor_fn(
    node_type: NodeType,
    write_tensor_by_layer_index: dict[LayerIndex, torch.Tensor] | dict[int, torch.Tensor] | None,
    layer_precedes_direction_layer_fn: Callable[[LayerIndex], bool],
) -> Callable[[torch.Tensor, torch.Tensor, LayerIndex, PassType], torch.Tensor]:
    """
    To convert an "activation" tensor to a "projection of write to direction" tensor, we need to convert the
    "activation" to the residual stream basis (using a write tensor) if it is not already, and then project to
    the direction of interest. This function constructs the appropriate tensor operation to perform this projection
    based on the node_type of the activation tensor (passed as an argument). The write_tensor_by_layer_index argument
    defines the conversion to the residual stream basis, and the layer_precedes_direction_layer_fn argument is assumed
    to return True iff the derived scalar at a given layer index is upstream of the direction of interest.
    """
    match node_type:
        case NodeType.RESIDUAL_STREAM_CHANNEL:
            assert write_tensor_by_layer_index is None

            def inner_product_with_residual(
                residual: torch.Tensor,
                direction: torch.Tensor,
                layer_index: LayerIndex,
                pass_type: PassType,
            ) -> torch.Tensor:  # (num_sequence_tokens, 1)
                assert pass_type == PassType.FORWARD
                assert residual.ndim == 2
                assert residual.shape == direction.shape
                if layer_precedes_direction_layer_fn(layer_index):
                    return torch.einsum("td,td->t", residual, direction)[
                        :, None
                    ]  # sum over residual stream channels
                else:
                    return torch.zeros_like(residual[:, 0:1])

            return inner_product_with_residual
        case NodeType.MLP_NEURON | NodeType.AUTOENCODER_LATENT | NodeType.MLP_AUTOENCODER_LATENT | NodeType.ATTENTION_AUTOENCODER_LATENT:
            assert write_tensor_by_layer_index is not None

            def multiply_by_projection_to_direction(
                activations: torch.Tensor,
                direction: torch.Tensor,
                layer_index: LayerIndex,
                pass_type: PassType,
            ) -> (
                torch.Tensor
            ):  # (num_sequence_tokens, num_activations [i.e. num_neurons, num_latents])
                assert layer_index is not None
                assert layer_index in write_tensor_by_layer_index
                assert pass_type == PassType.FORWARD
                assert activations.ndim == direction.ndim == 2
                if layer_precedes_direction_layer_fn(layer_index):
                    write_projection = torch.einsum(
                        "ao,to->ta", write_tensor_by_layer_index[layer_index], direction
                    )
                    return activations * write_projection
                else:
                    return torch.zeros_like(activations)

            return multiply_by_projection_to_direction
        case NodeType.V_CHANNEL:
            assert write_tensor_by_layer_index is not None

            def attn_write_to_residual_direction_tensor_calculate_derived_scalar_fn(
                attn_weighted_values: torch.Tensor,
                direction: torch.Tensor,
                layer_index: LayerIndex,
                pass_type: PassType,
            ) -> (
                torch.Tensor
            ):  # (num_sequence_tokens, num_heads) or (num_sequence_tokens, num_attended_to_sequence_tokens, num_heads)
                assert layer_index is not None
                assert layer_index in write_tensor_by_layer_index
                assert pass_type == PassType.FORWARD
                # one or two token dimensions, one head dimension, one value channel dimension
                assert attn_weighted_values.ndim in {3, 4}
                if layer_precedes_direction_layer_fn(layer_index):
                    return compute_attn_write_to_residual_direction_from_attn_weighted_values(
                        attn_weighted_values=attn_weighted_values,
                        residual_direction=direction,
                        W_O=write_tensor_by_layer_index[layer_index],
                        pass_type=pass_type,
                    )  # TODO: consider splitting into two cases, once we have separate node_types
                else:
                    return torch.zeros_like(attn_weighted_values[..., 0])  # sum over last dimension

            return attn_write_to_residual_direction_tensor_calculate_derived_scalar_fn
        case _:
            raise NotImplementedError(
                f"make_write_to_direction_tensor_fn not implemented for {node_type=}"
            )

    To convert an "activation" tensor to a "projection of write to direction" tensor, we need to convert the
    "activation" to the residual stream basis (using a write tensor) if it is not already, and then project to
    the direction of interest. This function constructs the appropriate tensor operation to perform this projection
    based on the node_type of the activation tensor (passed as an argument). The write_tensor_by_layer_index argument
    defines the conversion to the residual stream basis, and the layer_precedes_direction_layer_fn argument is assumed
    to return True iff the derived scalar at a given layer index is upstream of the direction of interest.
    
neuron_explainer/activations/derived_scalars/direct_effects.py