code: def run_multiprocess_workflow(process_id=0, cluster_size=1):
      jax.distributed.initialize(
          coordinator_address="localhost:12321",
          num_processes=cluster_size,
          process_id=process_id,
      )
  
      log.basicConfig(
          format=f"PID {process_id}: %(message)s", level=log.INFO
      )
  
      print_devices(process_id=process_id)
  
      test_positional_sharding_workflow(process_id=process_id)
      test_named_sharding_workflow(process_id=process_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 7 for calling a global function with example usage: run_multiprocess_workflow(process_id, cluster_size) without return types