code: def build_input_pipeline(
      x_train, x_test, y_train, y_test, batch_size, valid_size
  ):
      """Build an Iterator switching between train and heldout data."""
  
      x_train = x_train.astype("float32")
      x_test = x_test.astype("float32")
  
      x_train /= 255
      x_test /= 255
  
      y_train = y_train.flatten()
      y_test = y_test.flatten()
  
      if FLAGS.subtract_pixel_mean:
          x_train_mean = np.mean(x_train, axis=0)
          x_train -= x_train_mean
          x_test -= x_train_mean
  
      print("x_train shape:" + str(x_train.shape))
      print(str(x_train.shape[0]) + " train samples")
      print(str(x_test.shape[0]) + " test samples")
  
      # Build an iterator over training batches.
      training_dataset = tf.data.Dataset.from_tensor_slices(
          (x_train, np.int32(y_train))
      )
      training_batches = (
          training_dataset.shuffle(50000, reshuffle_each_iteration=True)
          .repeat()
          .batch(batch_size)
      )
      training_iterator = tf.compat.v1.data.make_one_shot_iterator(
          training_batches
      )
  
      # Build a iterator over the heldout set with batch_size=heldout_size,
      # i.e., return the entire heldout set as a constant.
      heldout_dataset = tf.data.Dataset.from_tensor_slices(
          (x_test, np.int32(y_test))
      )
      heldout_batches = heldout_dataset.repeat().batch(valid_size)
      heldout_iterator = tf.compat.v1.data.make_one_shot_iterator(
          heldout_batches
      )
  
      # Combine these into a feedable iterator that can switch between training
      # and validation inputs.
      handle = tf.compat.v1.placeholder(tf.string, shape=[])
      feedable_iterator = tf.compat.v1.data.Iterator.from_string_handle(
          handle,
          training_batches.output_types,
          training_batches.output_shapes,
      )
      images, labels = feedable_iterator.get_next()
  
      return images, labels, handle, training_iterator, heldout_iterator
  
file_path: probability/tensorflow_probability/examples/cifar10_bnn.py
imports: import os
  import warnings
  from absl import flags
  import matplotlib
  import numpy
  import tensorflow.compat.v1
  import tensorflow_probability
  from tensorflow_probability.examples.models.bayesian_resnet import bayesian_resnet
  from tensorflow_probability.examples.models.bayesian_vgg import bayesian_vgg
  from tensorflow_probability.python.internal import tf_keras
label_desc: Use image node_id 1 for calling a global function with example usage: build_input_pipeline(x_train, x_test, y_train, y_test, batch_size, valid_size) and returns: images,  labels,  handle,  training_iterator,  heldout_iterator