class_docstr: Adds 2-dimensional positional embeddings to the inputs.
  
  Attributes:
    rescale_from: tuple; If not None, embeddings are rescaled from this shape.
    posemb_init: Positional embedding initializer.
class_docstr_tok: ["Adds","2-dimensional","positional","embeddings","to","the","inputs",".","Attributes",":","rescale_from",":","tuple",";","If","not","None",",","embeddings","are","rescaled","from","this","shape",".","posemb_init",":","Positional","embedding","initializer","."]
docstr: Applies Add2DPositionEmbedding module.
  
  Args:
    inputs: nd-arrary; Input data.
  
  Returns:
    Output: `(bs, h, w, c)`.
docstr_tok: ["Applies","Add2DPositionEmbedding","module",".","Args",":","inputs",":","nd-arrary",";","Input","data",".","Returns",":","Output",":","`","(","bs",",","h",",","w",",","c",")","`","."]
code: def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
      """Applies Add2DPositionEmbedding module.
  
      Args:
        inputs: nd-arrary; Input data.
  
      Returns:
        Output: `(bs, h, w, c)`.
      """
      assert inputs.ndim == 4, (
          "Number of dimensions should be 4,"
          " but it is: %d" % inputs.ndim
      )
      _, h, w, c = inputs.shape
      embedding_h, embedding_w = h, w
      if self.rescale_from:  # `[h, w, c]`
          embedding_h, embedding_w = (
              self.rescale_from[0],
              self.rescale_from[1],
          )
  
      row_pos_embed = self.param(
          "row_pos_embedding", self.posemb_init, (embedding_w, c // 2)
      )
      col_pos_embed = self.param(
          "col_pos_embedding", self.posemb_init, (embedding_h, c // 2)
      )
      # To `[h, w, c//2]`.
      x_pos_emb = jnp.tile(
          jnp.expand_dims(row_pos_embed, axis=0), (embedding_h, 1, 1)
      )
      # To `[h, w, c//2]`.
      y_pos_emb = jnp.tile(
          jnp.expand_dims(col_pos_embed, axis=1), (1, embedding_w, 1)
      )
      # To `[h, w, c]`.
      pos = jnp.concatenate((x_pos_emb, y_pos_emb), axis=-1)
  
      if w != embedding_w or h != embedding_h:
          pos = jax.image.resize(
              pos, (h, w, c), method="bilinear", antialias=False
          )
  
      # To `[1, h, w, c]`.
      pos = jnp.expand_dims(pos, axis=0)
  
      return inputs + pos
  
imports: import functools
  from typing import Callable, Optional, Sequence, Tuple
  import flax.linen
  import jax
  import jax.numpy
  import numpy
  from scenic.model_lib.layers import nn_layers
label_desc: Use image node_id 1 for calling the Add2DPositionEmbedding obj's underlying member method code with example usage: obj.__call__(inputs) and returns: unknown