code: def resize_posemb(posemb, target_size):
      """Resizes position embeddings to new resolution."""
      if target_size == posemb.shape[-2]:
          return posemb
      posemb = jax.device_get(posemb)
      ndim = posemb.ndim
      if ndim == 3:
          posemb = posemb[0]
  
      gs_old = int(np.sqrt(posemb.shape[0]))
      gs_new = int(np.sqrt(target_size))
  
      posemb_tok = None
      if gs_old**2 == posemb.shape[0]:  # No CLS token.
          posemb_grid = posemb
      elif gs_old**2 == posemb.shape[0] - 1:  # Prepended CLS token.
          posemb_tok, posemb_grid = posemb[:1], posemb[1:]
      else:
          raise ValueError(
              "Posemb shape must be a perfect square (maybe with CLS token), but "
              f"got posemb of shape {posemb.shape}."
          )
  
      logging.info("Posemb: grid-size from %s to %s", gs_old, gs_new)
      posemb_grid = posemb_grid.reshape(gs_old, gs_old, -1)
  
      zoom = (gs_new / gs_old, gs_new / gs_old, 1)
      posemb_grid = scipy.ndimage.zoom(posemb_grid, zoom, order=1)
      posemb = posemb_grid.reshape(gs_new * gs_new, -1)
      if posemb_tok is not None:
          posemb = np.concatenate([posemb_tok, posemb], axis=0)
  
      return jnp.array(posemb[np.newaxis] if ndim == 3 else posemb)
  
imports: from typing import Any, Dict, List, Optional, Sequence, Tuple, Type, Union
  from absl import logging
  from clu import preprocess_spec
  from flax import core
  import jax
  import jax.numpy
  import numpy
  from scenic.train_lib import train_utils
  import scipy
label_desc: Use image node_id 1 for calling a global function with example usage: resize_posemb(posemb, target_size) and returns: jnp, posemb