code: def mnist_pipeline(data_path, random_shuffle):
      jpegs, labels = fn.readers.caffe2(
          path=data_path,
          random_shuffle=random_shuffle,
          name="mnist_caffe2_reader",
      )
      images = fn.decoders.image(
          jpegs, device="mixed", output_type=types.GRAY
      )
      images = fn.crop_mirror_normalize(
          images, dtype=types.FLOAT, std=[255.0], output_layout="HWC"
      )
  
      labels = labels.gpu()
      labels = fn.reshape(labels, shape=[])
  
      return images, labels
  
file_path: DALI/docs/examples/frameworks/jax/pax_examples/dali_pax_input.py
imports: import os
  from praxis import base_input
  from nvidia.dali.plugin import jax
  from nvidia.dali import pipeline_def
  import nvidia.dali.fn
  import nvidia.dali.types
label_desc: Use image node_id 1 for calling a global function with example usage: mnist_pipeline(data_path, random_shuffle) and returns: images,  labels