code: def tpu_system_init_helper(
      task_id,
      num_tasks,
      num_devices,
      use_tfrt_host_runtime=True,
      use_megacore=False,
  ):
      """A helper function to initialize multi-client tpu system."""
  
      @def_function.function
      def _tpu_init_fn():
          return gen_dtensor_ops.configure_and_initialize_global_tpu(
              use_tfrt_host_runtime=use_tfrt_host_runtime
          )
  
      @def_function.function
      def _set_global_tpu_array_fn(topology_proto):
          gen_dtensor_ops.d_tensor_set_global_tpu_array(topology_proto)
  
      with ops.device(
          "/job:" + config.full_job_name() + "/device:TPU_SYSTEM:0"
      ):  # pylint: disable=protected-access
          my_core_ids = _tpu_init_fn()
  
      if use_megacore:
          logging.info("Using TPU megacore")
          my_core_ids = my_core_ids * 2
      logging.info("TPU core IDs: %s", my_core_ids)
  
      # `my_core_ids` contains the IDs of TPU cores attached to this host.
      #
      # To generate correct and efficient XLA AllReduce group assignment, we must
      # merge these arrays from all hosts and broadcast the result back to all
      # hosts, so all hosts can use these mappings in their MLIR passes.
      #
      # This is essentially doing what WaitForDistributedTpuOp and
      # SetGlobalTPUArrayOp do, in our multi-client environment.
      num_devices_per_task = int(num_devices / num_tasks)
  
      # Create a one-time use mesh and layout just for merging core IDs.
      mesh = layout_lib.Mesh(
          [_MESH_DIM_X],
          *_create_device_array(
              (num_devices,), _TPU_DEVICE_TYPE, config.client_id()
          ),
      )
      layout = layout_lib.Layout(
          [_MESH_DIM_X, layout_lib.UNSHARDED], mesh
      )
      device = dtensor_device.DTensorDevice(meshes=[mesh])
      logging.info(
          "TPU core locations: %s",
          device.tpu_core_ids_to_locations(my_core_ids),
      )
  
      # At this point, we don't know which cores are attached to other hosts.
      # The core ID mappings in the runtime haven't been set yet.
      #
      # The core ID merging AllReduce below is carefully written so it works
      # without needing correct core mappings to be set in the runtime. We will
      # use this AllReduce's result to set the core ID mappings, and all future
      # user-initiated AllReduces will use the mappings.
      #
      # The runtime is hard-coded to ignore core ID mappings on this AllReduce.
      all_core_ids = np.zeros([num_devices], dtype=np.int32)
      for i in range(len(my_core_ids)):
          all_core_ids[
              task_id * num_devices_per_task + i
          ] = my_core_ids[i]
  
      # Only one local device gets a valid input. To give an example, assume we have
      # 2 tasks and each of them has 8 local devices, then `all_core_ids` in task 0
      # will have 8 tensors, where 1 of them may have its value as
      # [0,1,2,3,4,5,6,7,0,0,0,0,0,0,0,0] and the other tensors are all zeros. For
      # task 1, the case may be one with [0,0,0,0,0,0,0,0,8,9,10,11,12,13,14,15]
      # and other 7 are all zeros.
      all_core_ids = constant_op.constant([all_core_ids])
      zeros = array_ops.zeros_like(all_core_ids)
      all_core_ids = [all_core_ids] + [zeros] * (
          num_devices_per_task - 1
      )
  
      # All devices on all hosts participate in one AllReduce, whose result will be
      # core IDs arranged by task-device ordinals. For the above example, the result
      # will be [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15].
      with ops.device(device.name):
          all_core_ids = device.pack(all_core_ids, layout)
          all_core_ids = math_ops.reduce_sum(all_core_ids, axis=[0])
          unpacked_all_tpu_ids = device.unpack(all_core_ids)
  
      all_core_ids = list(unpacked_all_tpu_ids[0].numpy())
      logging.info("All TPU core IDs: %s", all_core_ids)
  
      # Set the default core ID mappings in the runtime for legacy code and tests.
      #
      # Legacy code and tests create TPU meshes directly without using the
      # `create_tpu_mesh` function below. Those meshes have global device IDs
      # equal to TF task-device ordinals. The `all_core_ids` array happens to
      # arrange core IDs by TF task-device ordinals. Using this array on those
      # meshes guarantee correct although inefficient results.
      device.set_tpu_core_ids("", all_core_ids)
  
      # Remember enough global, immutable information to be able to build any ring
      # we want prescribed by `create_tpu_mesh` in the future.
      global _all_core_ids
      _all_core_ids = all_core_ids
  
      all_core_locations = device.tpu_core_ids_to_locations(
          all_core_ids
      )
      all_core_locations = [
          _CoreLocation(l[0], l[1], l[2], l[3])
          for l in all_core_locations
      ]
      global _all_core_locations
      _all_core_locations = all_core_locations
      logging.info("All TPU core locations: %s", all_core_locations)
  
      tpu_topology = _create_tpu_topology(
          all_core_locations, num_tasks, num_devices_per_task
      )
  
      _set_global_tpu_array_fn(tpu_topology.serialized())
      return tpu_topology, device
  
imports: import functools
  import time
  from typing import List, Optional, Dict
  import numpy
  from tensorflow.dtensor.python import config
  from tensorflow.dtensor.python import dtensor_device
  from tensorflow.dtensor.python import gen_dtensor_ops
  from tensorflow.dtensor.python import layout
  from tensorflow.python.eager import context
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.platform import tf_logging
  from tensorflow.python.tpu import topology
  from tensorflow.python.util.tf_export import tf_export
label_desc: Use image node_id 4 for calling a global function with example usage: tpu_system_init_helper(task_id, num_tasks, num_devices, use_tfrt_host_runtime, use_megacore) and returns: tpu_topology,  device, gen_dtensor_ops