code: def run_distributed_sharing_test(sharding, process_id):
      dali_local_shard = dax.integration._to_jax_array(
          get_dali_tensor_gpu(process_id, (1), np.int32, 0)
      )
  
      # Note: we pass only one local shard but the array virtually
      # combines all shards together
      dali_sharded_array = jax.make_array_from_single_device_arrays(
          shape=(2,), sharding=sharding, arrays=[dali_local_shard]
      )
  
      # This array should be backed only by one device buffer that holds
      # local part of the data. This buffer should be on the local device.
      assert len(dali_sharded_array.device_buffers) == 1
      assert dali_sharded_array.device_buffer == jnp.array([process_id])
      assert (
          dali_sharded_array.device_buffer.device()
          == jax.local_devices()[0]
      )
      assert (
          dali_sharded_array.device_buffer.device()
          == jax.devices()[process_id]
      )
  
file_path: DALI/dali/test/python/jax_plugin/jax_server.py
imports: import jax
  import jax.numpy
  import numpy
  import logging
  from test_integration import get_dali_tensor_gpu
  from jax.sharding import PositionalSharding, NamedSharding, PartitionSpec, Mesh
  import nvidia.dali.plugin.jax
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