neuron_explainer/activations/derived_scalars/raw_activations.py [226:234]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        activations: torch.Tensor,
        layer_index: LayerIndex,
        pass_type: PassType,
    ) -> torch.Tensor:
        assert pass_type == PassType.FORWARD, "Backward pass not defined for write"
        assert (
            layer_index in write_tensor_by_layer_index
        ), f"{layer_index=} not in {write_tensor_by_layer_index.keys()=}"
        return torch.einsum(
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



neuron_explainer/activations/derived_scalars/raw_activations.py [273:281]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        activations: torch.Tensor,
        layer_index: LayerIndex,
        pass_type: PassType,
    ) -> torch.Tensor:
        assert pass_type == PassType.FORWARD, "Backward pass not defined for write"
        assert (
            layer_index in write_tensor_by_layer_index
        ), f"{layer_index=} not in {write_tensor_by_layer_index.keys()=}"
        return torch.einsum(
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -