File size: 1,625 Bytes
a2ffd07 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | 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
@contextmanager
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) |