code: def upscale_nn(x: JaxArray, scale: int = 2) -> JaxArray:
      """Nearest neighbor upscale for image batches of shape (N, C, H, W).
  
      Args:
          x: input tensor of shape (N, C, H, W).
          scale: integer scaling factor.
  
      Returns:
          Output tensor of shape (N, C, H * scale, W * scale).
      """
      s = x.shape
      x = x.reshape(s[:2] + (s[2], 1, s[3], 1))
      x = jn.tile(x, (1, 1, 1, scale, 1, scale))
      return x.reshape(s[:2] + (scale * s[2], scale * s[3]))
  
file_path: objax/objax/functional/core/ops.py
imports: from typing import Union, Tuple
  import jax.nn
  from jax import numpy, lax
  from objax import util
  from objax.constants import Interpolate
  from objax.typing import JaxArray
label_desc: Use image node_id 4 for calling a global function with example usage: upscale_nn(x, scale) and returns: x