[WIP] Enable training SAEs on vision models

[luciaquirke]

• Update input flattening to support hidden tensors with more batch dimensions
• Enable specifying dummy inputs - the default transformers dummy input is of shape [3, 5] and uses the 'input_ids' key
• Improve logging when hookpoints aren't specified correctly

TODO:

• Enable training from command line (replace tokenizer with image processor)
• Decide whether to support image data in MemmapDataset
• Check if we can detect vision models and produce correct dummy input (key value, image shape) in the init

Maybe like

processor = AutoProcessor.from_pretrained(args.model, token=args.hf_token)
target_column = "pixel_values" if isinstance(processor, BaseImageProcessor) else "input_ids"

process.size.shortest_edge exists but no channel count