code: def upsample_2d(
      x: JaxArray,
      scale: Union[Tuple[int, int], int],
      method: Union[Interpolate, str] = Interpolate.BILINEAR,
  ) -> JaxArray:
      """Function to upscale 2D images.
  
      Args:
          x: input tensor.
          scale: int or tuple scaling factor
          method: str or UpSample interpolation methods e.g. ['bilinear', 'nearest'].
  
      Returns:
          upscaled 2d image tensor
      """
      s = x.shape
      assert (
          len(s) == 4
      ), f"{s} must have 4 dimensions to be upsampled, or you can try interpolate function."
      scale = util.to_tuple(scale, 2)
      y = jax.image.resize(
          x.transpose([0, 2, 3, 1]),
          shape=(s[0], s[2] * scale[0], s[3] * scale[1], s[1]),
          method=util.to_interpolate(method),
      )
      return y.transpose([0, 3, 1, 2])
  
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 3 for calling a global function with example usage: upsample_2d(x, scale, method) and returns: y