code: def _safe_flatten_dict(
      dct: dict[Any, Any]
  ) -> tuple[list[Any], tree_util.PyTreeDef]:
      # We avoid comparison between keys by just using the original order
      keys, values = [], []
      for key, value in dct.items():
          try:
              tree_util.tree_leaves(value)
          except:
              # If flattening fails, we substitute a sentinel object.
              value = cant_flatten
          keys.append(key)
          values.append(value)
      return tree_util.tree_flatten(_DictWrapper(keys, values))
  
imports: from __future__ import annotations
  from collections.abc import Hashable
  import dataclasses
  import inspect
  import threading
  from typing import Any, Protocol
  import jax
  from jax import tree_util
  from jax._src import core
  from jax._src import debugging
  from jax._src import traceback_util
  from jax._src import util
label_desc: Use image node_id 1 for calling a global function with example usage: _safe_flatten_dict(dct) and returns: tree_util