code: def max_pool(x_input, pool_size):
      """max_pool downsamples a feature map by 2X."""
      return tf.nn.max_pool(
          x_input,
          ksize=[1, pool_size, pool_size, 1],
          strides=[1, pool_size, pool_size, 1],
          padding="SAME",
      )
  
file_path: nni/examples/trials/mnist-tfv1/mnist.py
imports: import argparse
  import logging
  import math
  import tempfile
  import time
  import tensorflow
  from tensorflow.examples.tutorials.mnist import input_data
  import nni
label_desc: Use image node_id 2 for calling a global function with example usage: max_pool(x_input, pool_size) and returns: tf