From 53ebe2339ec9110541d2c47eb67ae1566ce3df7c Mon Sep 17 00:00:00 2001
From: SouthOcean <92636795+Xu-Jiaqi@users.noreply.github.com>
Date: Wed, 4 Mar 2026 19:24:32 +0800
Subject: [PATCH 1/3] Refactor get_local to handle multiple variable names

---
 visualizer/visualizer.py | 24 +++++++++++++-----------
 1 file changed, 13 insertions(+), 11 deletions(-)

diff --git a/visualizer/visualizer.py b/visualizer/visualizer.py
index b7bef8c..b7cc639 100644
--- a/visualizer/visualizer.py
+++ b/visualizer/visualizer.py
@@ -4,22 +4,22 @@ class get_local(object):
     cache = {}
     is_activate = False
 
-    def __init__(self, varname):
-        self.varname = varname
+    def __init__(self, varnames):
+        self.varnames = [varnames] if isinstance(varnames, str) else varnames
 
     def __call__(self, func):
         if not type(self).is_activate:
             return func
 
-        type(self).cache[func.__qualname__] = []
+        for varname in self.varnames:
+            type(self).cache[func.__qualname__+'.'+varname] = []
         c = Bytecode.from_code(func.__code__)
         extra_code = [
                          Instr('STORE_FAST', '_res'),
-                         Instr('LOAD_FAST', self.varname),
-                         Instr('STORE_FAST', '_value'),
+                         *sum([[Instr('LOAD_FAST', varname), Instr('STORE_FAST', f'_value_{varname}')] for varname in self.varnames], []),
                          Instr('LOAD_FAST', '_res'),
-                         Instr('LOAD_FAST', '_value'),
-                         Instr('BUILD_TUPLE', 2),
+                         *[Instr('LOAD_FAST', f'_value_{varname}') for varname in self.varnames],
+                         Instr('BUILD_TUPLE', 1+len(self.varnames)),
                          Instr('STORE_FAST', '_result_tuple'),
                          Instr('LOAD_FAST', '_result_tuple'),
                      ]
@@ -27,8 +27,10 @@ def __call__(self, func):
         func.__code__ = c.to_code()
 
         def wrapper(*args, **kwargs):
-            res, values = func(*args, **kwargs)
-            type(self).cache[func.__qualname__].append(values.detach().cpu().numpy())
+            res_values = func(*args, **kwargs)
+            res, values = res_values[0], res_values[1:]
+            for i,varname in enumerate(self.varnames):
+                type(self).cache[func.__qualname__+'_'+varname].append(values[i].detach().cpu())
             return res
         return wrapper
 
@@ -38,5 +40,5 @@ def clear(cls):
             cls.cache[key] = []
 
     @classmethod
-    def activate(cls):
-        cls.is_activate = True
+    def activate(cls, activate=True):
+        cls.is_activate = activate

From 8ac0ef01c820c0d79bc8eccc26f762d166456599 Mon Sep 17 00:00:00 2001
From: SouthOcean <92636795+Xu-Jiaqi@users.noreply.github.com>
Date: Wed, 4 Mar 2026 19:35:58 +0800
Subject: [PATCH 2/3] fix

---
 visualizer/visualizer.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/visualizer/visualizer.py b/visualizer/visualizer.py
index b7cc639..c8000b3 100644
--- a/visualizer/visualizer.py
+++ b/visualizer/visualizer.py
@@ -40,5 +40,5 @@ def clear(cls):
             cls.cache[key] = []
 
     @classmethod
-    def activate(cls, activate=True):
-        cls.is_activate = activate
+    def activate(cls):
+        cls.is_activate = True

From 65cda030e3ef2bd0c7b1e910c07c595775093e19 Mon Sep 17 00:00:00 2001
From: SouthOcean <92636795+Xu-Jiaqi@users.noreply.github.com>
Date: Wed, 4 Mar 2026 21:38:16 +0800
Subject: [PATCH 3/3] Fix cache key formatting in visualizer.py

---
 visualizer/visualizer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/visualizer/visualizer.py b/visualizer/visualizer.py
index c8000b3..d0e6cca 100644
--- a/visualizer/visualizer.py
+++ b/visualizer/visualizer.py
@@ -30,7 +30,7 @@ def wrapper(*args, **kwargs):
             res_values = func(*args, **kwargs)
             res, values = res_values[0], res_values[1:]
             for i,varname in enumerate(self.varnames):
-                type(self).cache[func.__qualname__+'_'+varname].append(values[i].detach().cpu())
+                type(self).cache[func.__qualname__+'.'+varname].append(values[i].detach().cpu())
             return res
         return wrapper