def vector_reshape_hook(self, input_tensor: torch.Tensor) -> torch.Tensor:
        """
        Some activations are >1-d tensors per token, but the torch JVP function expects
        1-d tensor -> 1-d tensor functions as input. This reshapes the input tensor to be 1-d,
        and stores the shape for inverting the reshape later.
        """
        self._input_activation_shape = input_tensor.shape
        return input_tensor.reshape(input_tensor.shape[0], -1)

        Some activations are >1-d tensors per token, but the torch JVP function expects
        1-d tensor -> 1-d tensor functions as input. This reshapes the input tensor to be 1-d,
        and stores the shape for inverting the reshape later.
        
neuron_explainer/activations/derived_scalars/reconstituter_class.py