code: def get_dali_tensor_gpu(
      value, shape, dtype, device_id=0
  ) -> TensorGPU:
      """Helper function to create DALI TensorGPU.
  
      Args:
          value : Value to fill the tensor with.
          shape : Shape for the tensor.
          dtype : Data type for the tensor.
  
      Returns:
          TensorGPU: DALI TensorGPU with provided shape and dtype filled
          with provided value.
      """
  
      @pipeline_def(num_threads=1, batch_size=1)
      def dali_pipeline():
          values = types.Constant(
              value=np.full(shape, value, dtype), device="gpu"
          )
  
          return values
  
      pipe = dali_pipeline(device_id=device_id)
      pipe.build()
      dali_output = pipe.run()
  
      return dali_output[0][0]
  
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 1 for calling a global function with example usage: get_dali_tensor_gpu(value, shape, dtype, device_id) and returns: dali_output, values