code: def resample(img, flow):
      # img, NCHW
      # flow, N2HW
      B = tf.shape(img)[0]
      c = tf.shape(img)[1]
      h = tf.shape(img)[2]
      w = tf.shape(img)[3]
      img_flat = tf.reshape(tf.transpose(img, [0, 2, 3, 1]), [-1, c])
  
      dx, dy = tf.unstack(flow, axis=1)
      xf, yf = tf.meshgrid(
          tf.cast(tf.range(w), tf.float32),
          tf.cast(tf.range(h), tf.float32),
      )
      xf = xf + dx
      yf = yf + dy
  
      alpha = tf.expand_dims(xf - tf.floor(xf), axis=-1)
      beta = tf.expand_dims(yf - tf.floor(yf), axis=-1)
  
      xL = tf.clip_by_value(
          tf.cast(tf.floor(xf), dtype=tf.int32), 0, w - 1
      )
      xR = tf.clip_by_value(
          tf.cast(tf.floor(xf) + 1, dtype=tf.int32), 0, w - 1
      )
      yT = tf.clip_by_value(
          tf.cast(tf.floor(yf), dtype=tf.int32), 0, h - 1
      )
      yB = tf.clip_by_value(
          tf.cast(tf.floor(yf) + 1, dtype=tf.int32), 0, h - 1
      )
  
      batch_ids = tf.tile(
          tf.expand_dims(tf.expand_dims(tf.range(B), axis=-1), axis=-1),
          [1, h, w],
      )
  
      def get(y, x):
          idx = tf.reshape(batch_ids * h * w + y * w + x, [-1])
          idx = tf.cast(idx, tf.int32)
          return tf.gather(img_flat, idx)
  
      val = tf.zeros_like(alpha)
      val += (
          (1 - alpha)
          * (1 - beta)
          * tf.reshape(get(yT, xL), [-1, h, w, c])
      )
      val += (
          (0 + alpha)
          * (1 - beta)
          * tf.reshape(get(yT, xR), [-1, h, w, c])
      )
      val += (
          (1 - alpha)
          * (0 + beta)
          * tf.reshape(get(yB, xL), [-1, h, w, c])
      )
      val += (
          (0 + alpha)
          * (0 + beta)
          * tf.reshape(get(yB, xR), [-1, h, w, c])
      )
  
      # we need to enforce the channel_dim known during compile-time here
      shp = img.shape.as_list()
      return tf.reshape(
          tf.transpose(val, [0, 3, 1, 2]), [-1, shp[1], h, w]
      )
  
imports: import tensorflow
  from tensorpack import ModelDesc, argscope, enable_argscope_for_module
label_desc: Use image node_id 4 for calling a global function with example usage: resample(img, flow) and returns: tf, tf