class_docstr: Layer for MultiScaleDeformableAttention.
class_docstr_tok: ["Layer","for","MultiScaleDeformableAttention","."]
code: def setup(self):
      # (nlevels, 2)
      arr_shapes = jnp.asarray(self.spatial_shapes)
      self.offset_norm = jnp.stack(
          [arr_shapes[:, 1], arr_shapes[:, 0]], -1
      )
      # (1, 1, 1, nlevels, 1, 2)
      self.offset_norm = self.offset_norm[None, None, None, :, None, :]
  
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 1 for calling the MultiScaleDeformableAttention obj's underlying member method code with example usage: obj.setup() without return types