code: def test_lax_workflow(process_id):
      array_from_dali = dax.integration._to_jax_array(
          get_dali_tensor_gpu(1, (1), np.int32)
      )
  
      assert (
          array_from_dali.device() == jax.local_devices()[0]
      ), "Array should be backed by the device local to current process."
  
      sum_across_devices = jax.pmap(
          lambda x: jax.lax.psum(x, "i"), axis_name="i"
      )(array_from_dali)
  
      assert sum_across_devices[0] == len(
          jax.devices()
      ), "Sum across devices should be equal to the number of devices as data per device = [1]"
  
      log.info("Passed lax workflow test")
  
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 3 for calling a global function with example usage: test_lax_workflow(process_id) without return types