def apply_layerwise_transform_fn_to_activations(
        self,
        layerwise_transform_fn: LayerwiseTransformFn,
        output_dst: DerivedScalarType,
        output_pass_type: PassType,
    ) -> "ActivationsAndMetadata":
        """layerwise_transform_fn is a Callable[[torch.Tensor, LayerIndex], torch.Tensor]. This function applies the transform_fn
        to the activations_by_layer_index at each layer index, and returns a new ActivationsAndMetadata
        object with the transformed activations."""

        def _layerwise_transform_activations_by_layer_index(
            activations_by_layer_index: dict[LayerIndex, torch.Tensor],
            layerwise_transform_fn: LayerwiseTransformFn,
        ) -> dict[LayerIndex, torch.Tensor]:
            return {
                layer_index: layerwise_transform_fn(activations, layer_index)
                for layer_index, activations in activations_by_layer_index.items()
            }

        return dataclasses.replace(
            self,
            activations_by_layer_index=_layerwise_transform_activations_by_layer_index(
                self.activations_by_layer_index,
                layerwise_transform_fn=layerwise_transform_fn,
            ),
            dst=output_dst,
            pass_type=output_pass_type,
        )
layerwise_transform_fn is a Callable[[torch.Tensor, LayerIndex], torch.Tensor]. This function applies the transform_fn
        to the activations_by_layer_index at each layer index, and returns a new ActivationsAndMetadata
        object with the transformed activations.
neuron_explainer/activations/derived_scalars/activations_and_metadata.py