code: def _flatten_grad_args(
      args,
  ) -> Tuple[list, list, List[bool], TreeSpec]:
      flat_args, spec = tree_flatten(args)
      grad_args = []
      other_args = []
      grad_flags = []
      for arg in flat_args:
          flag = _is_grad_tensor(arg)
          grad_flags.append(flag)
          if flag:
              grad_args.append(arg)
          else:
              other_args.append(arg)
      assert len(grad_args) > 0
      return grad_args, other_args, grad_flags, spec
  
imports: from abc import ABC, abstractmethod
  from contextlib import contextmanager
  from typing import Any, List, Tuple
  import torch
  from torch.utils._pytree import TreeSpec, tree_flatten, tree_unflatten
label_desc: Use image node_id 2 for calling a global function with example usage: _flatten_grad_args(args) and returns: grad_args,  other_args,  grad_flags,  spec