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/qa/TL3_JAX_multiprocess/jax_server.py
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 2 for calling a global function with example usage: print_devices(process_id) without return types