code: def _replicator_primary_device() -> snt_replicator.Replicator:
      # NOTE: The explicit device list is required since currently Replicator
      # only considers CPU and GPU devices. This means on TPU by default we only
      # mirror on the local CPU.
      for device_type in ("TPU", "GPU", "CPU"):
          devices = tf.config.experimental.list_logical_devices(
              device_type=device_type
          )
          if devices:
              devices = [d.name for d in devices]
              logging.info("Replicating over %s", devices)
              return snt_replicator.Replicator(devices=devices)
  
      assert False, "No TPU/GPU or CPU found"
  
imports: from typing import Callable, Sequence, Tuple
  import unittest
  from absl import logging
  from sonnet.src.distribute import replicator
  import tensorflow
label_desc: Use image node_id 1 for calling a global function with example usage: _replicator_primary_device() and returns: snt_replicator