code: def route_args(router, args, depth):
      routed_args = [(dict(), dict()) for _ in range(depth)]
      matched_keys = [key for key in args.keys() if key in router]
  
      for key in matched_keys:
          val = args[key]
          for depth, ((f_args, g_args), routes) in enumerate(
              zip(routed_args, router[key])
          ):
              new_f_args, new_g_args = map(
                  lambda route: ({key: val} if route else {}), routes
              )
              routed_args[depth] = (
                  {**f_args, **new_f_args},
                  {**g_args, **new_g_args},
              )
      return routed_args
  
imports: import torch
  import torch.nn
  from operator import itemgetter
  from torch.autograd.function import Function
  from torch.utils.checkpoint import get_device_states, set_device_states
label_desc: Use image node_id 1 for calling a global function with example usage: route_args(router, args, depth) and returns: routed_args