def apply_layerwise_transform_fn_to_output(
        self,
        layerwise_transform_fn: Callable[[torch.Tensor, LayerIndex, PassType], torch.Tensor],
        pass_type_to_transform: PassType,
        output_dst: DerivedScalarType,
    ) -> "ScalarDeriver":
        """Converts one ScalarDeriver to another, by applying a tensor -> tensor function to the output.
        The tensor -> tensor function takes a tensor, a layer index, and a pass type, and returns a tensor,
        so that it can depend on layer and pass type."""
        sub_scalar_sources = (DerivedScalarSource(self, pass_type=pass_type_to_transform),)

        def tensor_calculate_derived_scalar_fn(
            activation_data_tuple: tuple[torch.Tensor, ...],
            layer_index: LayerIndex,
            pass_type: PassType,
        ) -> torch.Tensor:
            assert len(activation_data_tuple) == 1
            return layerwise_transform_fn(activation_data_tuple[0], layer_index, pass_type)

        return dataclasses.replace(
            self,
            dst=output_dst,
            sub_scalar_sources=sub_scalar_sources,
            tensor_calculate_derived_scalar_fn=tensor_calculate_derived_scalar_fn,
            _specified_location_within_layer=self.location_within_layer,
        )
Converts one ScalarDeriver to another, by applying a tensor -> tensor function to the output.
        The tensor -> tensor function takes a tensor, a layer index, and a pass type, and returns a tensor,
        so that it can depend on layer and pass type.
neuron_explainer/activations/derived_scalars/scalar_deriver.py