code: def pipeline_with_variable_shape_output(batch_size):
      """Helper to create DALI pipelines that return GPU tensors with variable shape.
  
      Args:
          batch_size: Batch size for the pipeline.
      """
  
      def numpy_tensors(sample_info):
          tensors = [
              np.full((1, 5), sample_info.idx_in_epoch, dtype=np.int32),
              np.full((1, 3), sample_info.idx_in_epoch, dtype=np.int32),
              np.full((1, 2), sample_info.idx_in_epoch, dtype=np.int32),
              np.full((1, 4), sample_info.idx_in_epoch, dtype=np.int32),
          ]
          return tensors[sample_info.idx_in_epoch % len(tensors)]
  
      @pipeline_def(batch_size=batch_size, num_threads=4, device_id=0)
      def sequential_pipeline_def():
          data = fn.external_source(
              source=numpy_tensors,
              num_outputs=1,
              batch=False,
              dtype=types.INT32,
          )
          data = data[0].gpu()
          return data
  
      return sequential_pipeline_def()
  
file_path: DALI/dali/test/python/jax_plugin/utils.py
imports: import os
  import numpy
  from nvidia.dali import pipeline_def
  from nvidia.dali.backend import TensorGPU
  import nvidia.dali.fn
  import nvidia.dali.types
label_desc: Use image node_id 3 for calling a global function with example usage: pipeline_with_variable_shape_output(batch_size) and returns: sequential_pipeline_def, tensors, data