1) Fixed bug where output layers were duplicated, 2) Fixed metadata f…

…or containing_module_origin, 3) Fixed issue with saving buffer layers 4) Added a field that gives the argument names for the PyTorch function being called in the layer
johnmarktaylor91 · Jul 22, 2024 · 7611300 · 7611300
1 parent cf75e0f
commit 7611300
Show file tree

Hide file tree

Showing 2 changed files with 172 additions and 45 deletions.
diff --git a/torchlens/constants.py b/torchlens/constants.py
@@ -105,6 +105,8 @@
     "elapsed_time_total",
     "elapsed_time_function_calls",
     "elapsed_time_torchlens_logging",
+    # Lookup info
+    "func_argnames"
 ]
 
 TENSOR_LOG_ENTRY_FIELD_ORDER = [
@@ -158,6 +160,7 @@
     "func_call_stack",
     "func_time_elapsed",
     "func_rng_states",
+    "func_argnames",
     "num_func_args_total",
     "num_position_args",
     "num_keyword_args",