code: def reset_logical_devices(device_type, count):
      """Resets logical devices for CPU/GPU.
  
      Logical devices can only be instantiated once on a particular context. For
      now, context re-use is triggering some function duplication errors, so we
      reset the context on each call.
  
      Args:
        device_type: The device_type to reset.
        count: numbers of virtual device to reset to.
      """
      reset_context()
      devices = tf_config.list_physical_devices(device_type)
      if device_type.upper() not in ("CPU", "GPU"):
          raise ValueError(
              "resetting logical device for non-supported device type : "
              "%s" % device_type
          )
  
      if count < len(devices):
          devices = devices[:count]
          tf_config.set_visible_devices(
              devices, device_type=device_type.upper()
          )
  
      for i, device in enumerate(devices):
          n = (i + 1) * count // len(devices) - i * count // len(
              devices
          )
          assert n > 0  # guaranteed if count >= len(devices)
          configs = []
          for ordinal in range(n):
              if device_type.upper() == "GPU":
                  dev_config = context.LogicalDeviceConfiguration(
                      memory_limit=_DEFAULT_GPU_MEMORY_LIMIT,
                      experimental_device_ordinal=ordinal,
                  )
              else:
                  dev_config = context.LogicalDeviceConfiguration()
              configs.append(dev_config)
  
          tf_config.set_logical_device_configuration(device, configs)
  
imports: import collections
  import copy
  import itertools
  import json
  import os
  import typing
  from absl import flags
  from absl.testing import parameterized
  import numpy
  from tensorflow.dtensor.python import accelerator_util
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import config
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python import numpy_util
  from tensorflow.dtensor.python.config import is_gpu_present
  from tensorflow.dtensor.python.config import is_tpu_present
  from tensorflow.dtensor.python.config import preferred_device_type
  from tensorflow.dtensor.python.config import use_multi_device_mode
  from tensorflow.dtensor.python.tests import test_backend_util
  from tensorflow.dtensor.python.tests.test_backend_name import DTENSOR_TEST_UTIL_BACKEND
  from tensorflow.dtensor.python.tests.test_backend_name import DTensorTestUtilBackend
  from tensorflow.dtensor.python.tests.test_backend_util import DTensorTestBackendConfigurator
  from tensorflow.python.compat import v2_compat
  from tensorflow.python.eager import context
  from tensorflow.python.framework import config
  from tensorflow.python.framework import device
  from tensorflow.python.ops import resource_variable_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 6 for calling a global function with example usage: reset_logical_devices(device_type, count) without return types