code: def print_devices(process_id):
      log.info(
          f"Local devices = {jax.local_device_count()}, "
          f"global devices = {jax.device_count()}"
      )
  
      log.info("All devices: ")
      print_devices_details(jax.devices(), process_id)
  
      log.info("Local devices:")
      print_devices_details(jax.local_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 1 for calling a global function with example usage: print_devices(process_id) without return types