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 setup(self):
      self.conv = nn.Conv(
          self.in_channels,
          kernel_size=(3, 3),
          strides=(1, 1),
          padding=((1, 1), (1, 1)),
          dtype=self.dtype,
      )
  
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 1 for calling the FlaxUpsample2D obj's underlying member method code with example usage: obj.setup() without return types