code: def dtensor_reduce(strategy, reduce_op, value, axis):
      """Implement dtensor based strategy.reduce()."""
      # Due to the limitation of using scalar in DTensor (e.g. the rank 0 tensor
      # loss the batch shard information), we need to override the default
      # reduce in addition to the strategy.extend._reduce_to()
      # Most of the logic here is a mimic of the parent class, except for how
      # mean and sum are calculated in a global context.
      distribute_lib._require_cross_replica_or_default_context_extended(  # pylint: disable=protected-access
          strategy.extended
      )
      if isinstance(reduce_op, str):
          reduce_op = reduce_util.ReduceOp(reduce_op.upper())
  
      distributed_input = is_distributed_value(value)
      if not distributed_input and axis is None:
          # For any value that isn't distributed and doesn't need a reduction within
          # the replica.
          destinations = (
              device_util.current()
              or strategy.extended._default_device
              or "/device:CPU:0"  # pylint: disable=protected-access
          )
          devices = cross_device_ops_lib.get_devices_from(destinations)
          with ops.device(devices[0]):
              return array_ops.identity(
                  cross_device_ops_lib.reduce_non_distributed_value(
                      reduce_op,
                      value,
                      destinations,
                      strategy.num_replicas_in_sync,
                  )
              )
  
      value = convert_inputs_to_dtensor(
          value, strategy._mesh
      )  # pylint: disable=protected-access
      # At this point, the value is a DTensor instance now.
      # There will be a final reduction step cross replica. In order to maintain
      # the shape of each local replica, we need to add a new dim to the front.
      # E.g. 2 replica with local shape as (4, 5, 6), the global tensor shape
      # should be (8, 5, 6), we will reshape into (2, 4, 5, 6) and then do a
      # reduction on axis 0.
      if reduce_op == reduce_util.ReduceOp.MEAN:
          reduce_op = math_ops.reduce_mean
      else:
          reduce_op = math_ops.reduce_sum
  
      # TODO(scottzhu): Make sure we handle dynamic/uneven shape in future.
      if d_api.fetch_layout(value).is_fully_replicated():
          # In case of fully mirrored dtensor, we only need to do one reduce, and
          # don't need to care about any per-replica logic.
          if axis is not None:
              value = reduce_op(value, axis=axis)
      else:
          new_shape = [strategy.num_replicas_in_sync, -1]
          if len(value.shape) > 1:
              new_shape.extend(array_ops.shape(value)[1:])
          value = array_ops.reshape(value, new_shape)
          if axis is not None:
              # we do a reduce_sum/mean within each of the replica when axis is not
              # None. Add 1 to the axis since there is a new dim added by reshape in
              # front.
              value = reduce_op(value, axis=axis + 1)
          value = reduce_op(value, axis=0)
  
      # Note that we return a DTensor instance here, which should have the same
      # value as the original MirroredStrategy, but with a different type. User
      # might want a tf.Tensor for the status quo.
      return value
  
imports: from tensorflow.dtensor.python import accelerator_util
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import input_util
  from tensorflow.dtensor.python import layout
  from tensorflow.python.distribute import cross_device_ops
  from tensorflow.python.distribute import device_util
  from tensorflow.python.distribute import distribute_lib
  from tensorflow.python.distribute import reduce_util
  from tensorflow.python.distribute import values
  from tensorflow.python.eager import context
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_conversion_registry
  from tensorflow.python.framework import tensor_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import summary_ops_v2
label_desc: Use image node_id 7 for calling a global function with example usage: dtensor_reduce(strategy, reduce_op, value, axis) and returns: value, array_ops