code: def _callback_op_sharding(
      axis_context, sharding: SingleDeviceSharding | None
  ):
      if isinstance(axis_context, sharding_impls.SPMDAxisContext):
          # 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.
          if axis_context.manual_axes != frozenset(
              axis_context.mesh.axis_names
          ):
              raise NotImplementedError(
                  "callbacks are only supported in spmd computations when all mesh"
                  " axes are partitioned manually (no partial automatic sharding)."
              )
          if sharding is not None:
              raise NotImplementedError(
                  "callbacks do not support specifying sharding inside spmd"
                  " computations"
              )
          op_sharding = xc.OpSharding()
          op_sharding.type = xc.OpSharding.Type.MANUAL
          return op_sharding
  
      if isinstance(axis_context, sharding_impls.ShardingContext):
          if sharding is not None:
              if not isinstance(sharding, SingleDeviceSharding):
                  raise NotImplementedError(
                      "pure_callback only supports SingleDeviceSharding, but got"
                      f" {type(sharding)}"
                  )
              device = next(iter(sharding.device_set))
              device_assignment = axis_context.device_assignment
              if device_assignment is None:
                  raise AssertionError(
                      "Please file a bug at https://github.com/google/jax/issues"
                  )
              try:
                  device_index = device_assignment.index(device)
              except IndexError as e:
                  raise ValueError(
                      "Sharding provided to pure_callback specifies a device"
                      f" {device} that is not in the device assignment"
                      f" ({device_assignment})"
                  ) from e
          else:
              device_index = 0
  
          # 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).
          op_sharding = xc.OpSharding()
          op_sharding.type = xc.OpSharding.Type.MAXIMAL
          op_sharding.tile_assignment_dimensions = [1]
          op_sharding.tile_assignment_devices = [device_index]
          return op_sharding
  
      # When there's no SPMD partitioning going on, don't annotate a sharding.
      return None
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import functools
  from typing import Any, Callable
  import numpy
  from jax._src import core
  from jax._src import dispatch
  from jax._src import dtypes
  from jax._src import effects
  from jax._src import sharding_impls
  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.lib import xla_client
  from jax._src.lax.control_flow.loops import map
  from jax._src.sharding_impls import SingleDeviceSharding
label_desc: Use image node_id 6 for calling a global function with example usage: _callback_op_sharding(axis_context, sharding) and returns: None, op_sharding, op_sharding