code: def tree_flatten_with_names(pytree, is_leaf, path="", to_id=id):
      id_to_name = {}
      if getattr(pytree, "items", None):
          for k, v in pytree.items():
              k_path = f"{path}/{k}"
              if is_leaf(v):
                  id_to_name[to_id(v)] = k_path
              else:
                  id_to_name = {
                      **id_to_name,
                      **tree_flatten_with_names(
                          v, is_leaf=is_leaf, path=k_path
                      ),
                  }
      elif getattr(pytree, "__getitem__", None):
          for v in pytree:
              if is_leaf(v):
                  id_to_name[to_id(v)] = path
              else:
                  id_to_name = {
                      **id_to_name,
                      **tree_flatten_with_names(
                          v, is_leaf=is_leaf, path=path
                      ),
                  }
      else:
          id_to_name[to_id(pytree)] = path
      return id_to_name
  
imports: import functools
  import io
  import json
  import time
  import jax
  import jax.numpy
  import numpy
  import multiprocessing
  import ray
  from smart_open import open
  from mesh_transformer.util import head_print
label_desc: Use image node_id 12 for calling a global function with example usage: tree_flatten_with_names(pytree, is_leaf, path, to_id) and returns: id_to_name