code: def create_device_ids_array(shape):
      device_count = np.prod(shape)
      return np.arange(device_count).reshape(shape)
  
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 2 for calling a global function with example usage: create_device_ids_array(shape) and returns: np