code: def augment(
      x, shift=4
  ):  # Shift all images in the batch by up to "shift" pixels in any direction.
      x_pad = np.pad(
          x, [[0, 0], [0, 0], [shift, shift], [shift, shift]]
      )
      rx, ry = np.random.randint(0, shift, size=2)
      return x_pad[:, :, rx : rx + 28, ry : ry + 28]
  
file_path: objax/examples/image_classification/mnist_cnn.py
imports: import os
  import numpy
  import tensorflow_datasets
  from tqdm import trange
  import objax
  from objax.util import EasyDict
label_desc: Use image node_id 2 for calling a global function with example usage: augment(x, shift) and returns: x_pad