code: def tree_flatten_hf(pytree: Any) -> Tuple[List[Any], TreeSpec]:
      """Flattens a pytree into a list of values an a TreeSpec that can be used
      to reconstruct the pytree.
      """
      if isinstance(pytree, OrderedDict):
          node_type = OrderedDict
          flatten_fn = SUPPORTED_NODES[node_type].flatten_fn
          child_pytrees, context = flatten_fn(pytree)
  
          # Recursively flatten the children
          result: List[Any] = []
          children_specs: List["TreeSpec"] = []
          for child in child_pytrees:
              flat, child_spec = tree_flatten_hf(child)
              result += flat
              children_specs.append(child_spec)
          return result, TreeSpec(node_type, context, children_specs)
      else:
          result, tree_spec = tree_flatten(pytree)
          return result, tree_spec
  
file_path: colossalai/colossalai/pipeline/schedule/_utils.py
imports: from collections import OrderedDict
  from typing import Any, List, Optional, Tuple
  import torch
  import torch.cuda
  from torch.nn import Module
  from torch.utils._pytree import SUPPORTED_NODES, TreeSpec, _register_pytree_node, tree_flatten, tree_map, tree_unflatten
label_desc: Use image node_id 4 for calling a global function with example usage: tree_flatten_hf(pytree) and returns: result,  TreeSpec, result,  tree_spec