code: def posemb_sincos_2d(
      h, w, width, temperature=10_000.0, dtype=jnp.float32
  ):
      """Follows the MoCo v3 logic."""
      y, x = jnp.mgrid[:h, :w]
  
      assert width % 4 == 0, "Width must be mult of 4 for sincos posemb"
      omega = jnp.arange(width // 4) / (width // 4 - 1)
      omega = 1.0 / (temperature**omega)
      y = jnp.einsum("m,d->md", y.flatten(), omega)
      x = jnp.einsum("m,d->md", x.flatten(), omega)
      pe = jnp.concatenate(
          [jnp.sin(x), jnp.cos(x), jnp.sin(y), jnp.cos(y)], axis=1
      )
      return jnp.asarray(pe, dtype)[None, :, :]
  
imports: import functools
  from typing import Any, Optional
  from absl import logging
  import flax
  import flax.linen
  from flax.training import common_utils
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.base_models import multilabel_classification_model
  from scenic.model_lib.layers import nn_layers
  import scipy
label_desc: Use image node_id 1 for calling a global function with example usage: posemb_sincos_2d(h, w, width, temperature, dtype) and returns: jnp