class_docstr: Layer for MultiScaleDeformableAttention.
class_docstr_tok: ["Layer","for","MultiScaleDeformableAttention","."]
docstr: Initializes deformable attention sampling offsets.
docstr_tok: ["Initializes","deformable","attention","sampling","offsets","."]
code: def pos_grid_init(
      self, key: jax.Array, shape: Shape, dtype: jnp.dtype
  ) -> Array:
      """Initializes deformable attention sampling offsets."""
      del key, shape
      thetas = (
          jnp.arange(self.num_heads, dtype=dtype)
          * (2 * math.pi)
          / self.num_heads
      )
      grid_init = jnp.stack([jnp.cos(thetas), jnp.sin(thetas)], -1)
      denom = jnp.max(jnp.abs(grid_init), axis=-1, keepdims=True)
      grid_init = grid_init / denom
      grid_init = grid_init.reshape(self.num_heads, 1, 1, 2)
      grid_init = jnp.tile(
          grid_init, (1, self.num_levels, self.num_points, 1)
      )
      for i in range(self.num_points):
          grid_i = grid_init[:, :, i, :] * (i + 1)
          grid_init = grid_init.at[:, :, i, :].set(grid_i)
      return grid_init
  
file_path: scenic/scenic/projects/baselines/deformable_detr/attention.py
imports: import functools
  import math
  from typing import Callable, Sequence, Tuple, Union
  import flax.linen
  import jax
  from jax.core import Shape
  import jax.numpy
  import ml_collections
  import numpy
label_desc: Use image node_id 2 for calling the MultiScaleDeformableAttention obj's underlying member method code with example usage: obj.pos_grid_init(key, shape, dtype) and returns: grid_init