code: def run_sharding_test(sharding):
      # given
      dali_shard_0 = get_dali_tensor_gpu(0, (1), np.int32, 0)
      dali_shard_1 = get_dali_tensor_gpu(1, (1), np.int32, 1)
  
      shards = [
          dax.integration._to_jax_array(dali_shard_0),
          dax.integration._to_jax_array(dali_shard_1),
      ]
  
      assert shards[0].device() == jax.devices()[0]
      assert shards[1].device() == jax.devices()[1]
  
      # when
      dali_sharded_array = jax.make_array_from_single_device_arrays(
          shape=(2,), sharding=sharding, arrays=shards
      )
  
      # then
      jax_sharded_array = jax.device_put(jnp.arange(2), sharding)
  
      assert (dali_sharded_array == jax_sharded_array).all()
      assert (
          len(dali_sharded_array.device_buffers) == jax.device_count()
      )
  
      assert (
          dali_sharded_array.device_buffers[0].device()
          == jax.devices()[0]
      )
      assert (
          dali_sharded_array.device_buffers[1].device()
          == jax.devices()[1]
      )
  
imports: import jax
  import numpy
  import nvidia.dali.plugin.jax
  from nvidia.dali import pipeline_def
  import nvidia.dali.fn
  import nvidia.dali.types
  from nvidia.dali.plugin.jax import DALIGenericIterator, data_iterator
  from nvidia.dali.plugin.base_iterator import LastBatchPolicy
  from jax.sharding import PositionalSharding, NamedSharding, PartitionSpec, Mesh
  from jax.experimental import mesh_utils
  from utils import get_dali_tensor_gpu, iterator_function_def
  import jax.numpy
  import itertools
  from nose_utils import raises
label_desc: Use image node_id 4 for calling a global function with example usage: run_sharding_test(sharding) without return types