Refactor bytecode_to_tree.retrieve_task_root

kernel/modelverse_jit/bytecode_to_tree.py

@@ -34,10 +34,13 @@ def apply_intrinsic(intrinsic_function, named_args):
 def retrieve_task_root():
     """Creates an instruction that stores the task_root variable in a local."""
     return tree_ir.StoreLocalInstruction(
-        'task_root',
-        tree_ir.LoadIndexInstruction(
-            tree_ir.LoadLocalInstruction(jit_runtime.KWARGS_PARAMETER_NAME),
-            tree_ir.LiteralInstruction('task_root')))
+        'task_root', load_task_root())
+
+def load_task_root():
+    """Creates an instruction that loads the task_root variable."""
+    return tree_ir.LoadIndexInstruction(
+        tree_ir.LoadLocalInstruction(jit_runtime.KWARGS_PARAMETER_NAME),
+        tree_ir.LiteralInstruction('task_root'))
 
 def load_kernel():
     """Creates an instruction that loads the Modelverse kernel."""