| from unittest.mock import patch | |
| from contextlib import contextmanager, ExitStack | |
| # 1. A helper to generate the wrapper closure | |
| # We pass 'original_func' here so it is "baked in" to the wrapper. | |
| def create_avg_wrapper(original_func): | |
| def wrapper(inputs, *args, **kwargs): | |
| # Your custom logic | |
| if len(inputs.shape) == 3: | |
| inputs = inputs.mean(dim=1, keepdim=True) | |
| elif len(inputs.shape) == 2: | |
| inputs = inputs.unsqueeze(1) | |
| else: | |
| raise | |
| # Call the CAPTURED original function | |
| return original_func(inputs, *args, **kwargs) | |
| return wrapper | |
| def vision_sae_wrapper(saes): | |
| # Ensure we can handle a single SAE or a list of them | |
| if not isinstance(saes, list): | |
| saes = [saes] | |
| # ExitStack is designed exactly for "I need N context managers" | |
| with ExitStack() as stack: | |
| for sae in saes: | |
| # 1. Capture the specific original method for THIS sae | |
| original_forward = sae.forward | |
| # 2. Create a wrapper specific to this sae (closing over original_forward) | |
| patched_forward = create_avg_wrapper(original_forward) | |
| # 3. Enter the patch context and add it to the stack | |
| stack.enter_context(patch.object(sae, 'forward', side_effect=patched_forward)) | |
| # Yield control back to your code | |
| yield | |
| # --- Usage Example --- | |
| # sae_list = [sae1, sae2, sae3] | |
| # with vision_sae_wrapper(sae_list): | |
| # # All SAEs are now patched | |
| # output1 = sae1(input_tensor) | |
| # output2 = sae2(input_tensor) |