class_docstr: Flax implementation of 2D Upsample layer
  
  Args:
      in_channels (`int`):
          Input channels
      dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):
          Parameters `dtype`
class_docstr_tok: ["Flax","implementation","of","2D","Upsample","layer","Args",":","in_channels","(","`","int","`",")",":","Input","channels","dtype","(",":","obj",":","`","jnp.dtype","`",",","*","optional","*",",","defaults","to","jnp.float32",")",":","Parameters","`","dtype","`"]
code: def __call__(self, hidden_states):
      batch, height, width, channels = hidden_states.shape
      hidden_states = jax.image.resize(
          hidden_states,
          shape=(batch, height * 2, width * 2, channels),
          method="nearest",
      )
      hidden_states = self.conv(hidden_states)
      return hidden_states
  
imports: import math
  from functools import partial
  from typing import Tuple
  import flax
  import flax.linen
  import jax
  import jax.numpy
  from flax.core.frozen_dict import FrozenDict
  from ..configuration_utils import ConfigMixin, flax_register_to_config
  from ..utils import BaseOutput
  from .modeling_flax_utils import FlaxModelMixin
label_desc: Use image node_id 2 for calling the FlaxUpsample2D obj's underlying member method code with example usage: obj.__call__(hidden_states) and returns: hidden_states