code: def upscale2x_nearest_neighbor(inputs):
      """Doubles image size by repeating every pixel 2x2 times.
  
      Args:
        inputs: nd-array: Inputs in shape of `[bs, height, width, channels]'
  
      Returns:
        Upscaled inputs, in shape of  `[bs, 2*height, 2*width, channels]'
      """
      input_channels = inputs.shape[-1]
      input_h, input_w = inputs.shape[1], inputs.shape[2]
      input_nchw = jnp.transpose(inputs, (0, 3, 1, 2))
      flat_input_shape = (-1, input_h, input_w, 1)
      flat_input = jnp.reshape(input_nchw, flat_input_shape)
  
      height_scale, width_scale = 2, 2
      resize_kernel = jnp.ones((height_scale, width_scale, 1, 1))
      strides = (height_scale, width_scale)
      flat_output = lax.conv_transpose(
          flat_input, resize_kernel, strides, padding="VALID"
      )
  
      output_nchw_shape = (
          -1,
          input_channels,
          height_scale * input_h,
          width_scale * input_w,
      )
      output_nchw = jnp.reshape(flat_output, output_nchw_shape)
      resized_x = jnp.transpose(
          output_nchw, (0, 2, 3, 1)
      )  # Output: nhwc.
      return resized_x
  
file_path: scenic/scenic/model_lib/layers/nn_ops.py
imports: from typing import Callable, Sequence
  import flax.linen
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 9 for calling a global function with example usage: upscale2x_nearest_neighbor(inputs) and returns: resized_x