def truncate_to_expected_shape(
    tensor: torch.Tensor,
    expected_shape: list[int | None],
) -> torch.Tensor:
    """
    This asserts that the tensor has the expected shape, and optionally truncates it to that shape.
    None in expected_shape means that dimension is not checked.
    """
    if expected_shape is None:
        return tensor
    for dim, real_size, expected_size in zip(
        range(len(expected_shape)), tensor.shape, expected_shape
    ):
        if expected_size is not None:
            assert (
                real_size >= expected_size
            ), f"Dimension {dim} of tensor has size {real_size} but expected size {expected_size}"
            tensor = tensor.narrow(dim, 0, expected_size)
    return tensor

    This asserts that the tensor has the expected shape, and optionally truncates it to that shape.
    None in expected_shape means that dimension is not checked.
    
neuron_explainer/activations/derived_scalars/raw_activations.py