code: def debug_callback_lowering(ctx, *args, effect, callback, **params):
      axis_context = ctx.module_context.axis_context
      if isinstance(
          axis_context, sharding_impls.SPMDAxisContext
      ) and set(axis_context.manual_axes) == set(
          axis_context.mesh.axis_names
      ):
          # If we have fully manual sharding during lowering, that means the JAX
          # program has per-device semantics, so we run the callback on each device.
          sharding = xc.OpSharding()
          sharding.type = xc.OpSharding.Type.MANUAL
      elif isinstance(
          axis_context,
          (
              sharding_impls.ShardingContext,
              sharding_impls.SPMDAxisContext,
          ),
      ):
          # If we have fully automatic sharding during lowering, that means the JAX
          # program has bulk array semantics, so we run the callback with a MAXIMAL
          # sharding and hence execute it only once on the full logical value).
          # If we have partially automatic sharding, we do this too... not sure why!
          sharding = xc.OpSharding()
          sharding.type = xc.OpSharding.Type.MAXIMAL
          sharding.tile_assignment_dimensions = [1]
          sharding.tile_assignment_devices = [0]
      else:
          # When there's no SPMD partitioning going on, don't annotate a sharding.
          sharding = None
  
      def _callback(*flat_args):
          return tuple(
              debug_callback_p.impl(
                  *flat_args, effect=effect, callback=callback, **params
              )
          )
  
      if effects.ordered_effects.contains(effect):
          token = ctx.tokens_in.get(effect)[0]
          result, token, _ = mlir.emit_python_callback(
              ctx,
              _callback,
              token,
              list(args),
              ctx.avals_in,
              ctx.avals_out,
              True,
          )
          ctx.set_tokens_out(mlir.TokenSet({effect: (token,)}))
      else:
          result, token, _ = mlir.emit_python_callback(
              ctx,
              _callback,
              None,
              list(args),
              ctx.avals_in,
              ctx.avals_out,
              True,
              sharding=sharding,
          )
      return result
  
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 6 for calling a global function with example usage: debug_callback_lowering(ctx) and returns: result, tuple