code: def run_distributed_sharing_test(sharding, process_id):
      log.info(f"Sharding: {sharding}")
  
      dali_local_shards = []
      for id, device in enumerate(jax.local_devices()):
          current_shard = dax.integration._to_jax_array(
              get_dali_tensor_gpu(process_id, (1), np.int32, id)
          )
  
          assert current_shard.device() == device
  
          dali_local_shards.append(current_shard)
  
      dali_sharded_array = jax.make_array_from_single_device_arrays(
          shape=(jax.device_count(),),
          sharding=sharding,
          arrays=dali_local_shards,
      )
  
      assert (
          len(dali_sharded_array.device_buffers)
          == jax.local_device_count()
      )
  
      for id, buffer in enumerate(dali_sharded_array.device_buffers):
          assert buffer == jnp.array([process_id])
          assert buffer.device() == jax.local_devices()[id]
  
imports: import jax
  import jax.numpy
  import numpy
  import logging
  import argparse
  from nvidia.dali import pipeline_def
  from nvidia.dali.backend import TensorGPU
  import nvidia.dali.types
  import nvidia.dali.plugin.jax
  from jax.sharding import Mesh
  from jax.sharding import PartitionSpec
  from jax.sharding import PositionalSharding
  from jax.sharding import NamedSharding
label_desc: Use image node_id 4 for calling a global function with example usage: run_distributed_sharing_test(sharding, process_id) without return types