code: def _debug_callback_partial_eval_custom(
      saveable, unks_in, inst_in, eqn
  ):
      # The default behavior for effectful primitives is to not stage them if
      # possible. For debug callback, we actually want it to be staged to
      # provide more information to the user. This rule bypasses partial_eval's
      # regular behavior to do that. Specifically, we will stage the callback
      # if:
      # 1) the policy says debug_callbacks are not saveable
      # 2) the policy says debug_callbacks are saveable BUT all of the input
      #    values are instantiated.
      # The purpose is to call back with as much information as possible while
      # avoiding unnecessarily staging out other values.
      if any(unks_in):
          # The usual case (if we have any unknowns, we need to stage it out)
          res = [v for v, inst in zip(eqn.invars, inst_in) if not inst]
          return None, eqn, [], [], res
      if saveable(
          debug_callback_p, *[v.aval for v in eqn.invars], **eqn.params
      ):
          # The policy is telling us we can save the debug callback.
          if all(inst_in):
              # If all of the inputs are instantiated, we also stage out the
              # debug_callback.
              return eqn, eqn, [], [], []
          else:
              # If any are not instantiated, we don't do any extra staging to avoid
              # affecting the computation.
              return eqn, None, [], [], []
      # If we can't save the debug callback (thanks to the policy) we listen to
      # the policy and stage out the debug callback.
      return eqn, eqn, [], [], []
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import functools
  import string
  import sys
  from typing import Any, Callable, Union
  import weakref
  import numpy
  import jax.numpy
  from jax import lax
  from jax._src import core
  from jax._src import effects
  from jax._src import mesh
  from jax._src import sharding_impls
  from jax._src import dispatch
  from jax._src import tree_util
  from jax._src import util
  from jax._src.interpreters import ad
  from jax._src.interpreters import batching
  from jax._src.interpreters import mlir
  from jax._src.interpreters import partial_eval
  from jax._src.lib import xla_client
  from jax._src.lib.mlir import ir
  from jax._src.lib.mlir.dialects import hlo
  from jax._src.sharding import Sharding
  from jax._src.sharding_impls import NamedSharding, parse_flatten_op_sharding
label_desc: Use image node_id 7 for calling a global function with example usage: _debug_callback_partial_eval_custom(saveable, unks_in, inst_in, eqn) and returns: eqn,  eqn,  list,  list,  list, None,  eqn,  list,  list,  res, eqn,  eqn,  list,  list,  list, eqn,  None,  list,  list,  list