code: def _deformable_attention_kernel(
      value: Tensor,
      value_spatial_shapes: list[tuple[int, int]],
      sampling_locations: Tensor,
      attention_weights: Tensor,
  ) -> Tensor:
      """Deformable Attention Kernel used in Deformable DETR.
  
      Described in https://arxiv.org/abs/2010.04159.
  
      Args:
          value: shape (N, Lv, n_head * C)
          value_spatial_shapes: [(H0, W0), (H1, W1), ...]
          sampling_locations: shape (N, Lq, n_head, n_levels, n_points, 2)
          attention_weights: shape (N, Lq, n_head, n_levels, n_points)
  
      Returns:
          output, shape (N, Lq, n_head * C)
      """
      bs, _, n_head, c = value.shape
      _, Len_q, _, n_levels, n_points, _ = sampling_locations.shape
  
      split_shape: list[int] = [h * w for h, w in value_spatial_shapes]
      value_list = value.split(split_shape, dim=1)
      sampling_grids = 2 * sampling_locations - 1
  
      sampling_value_list: list[Tensor] = []
      for level, (h, w) in enumerate(value_spatial_shapes):
          # N_, H_*W_, M_, D_ -> N_, H_*W_, M_*D_ -> N_, M_*D_, H_*W_ -> N_*M_, D_, H_, W_
          value_l_ = (
              value_list[level]
              .flatten(2)
              .permute(0, 2, 1)
              .reshape(bs * n_head, c, h, w)
          )
          # N_, Lq_, M_, P_, 2 -> N_, M_, Lq_, P_, 2 -> N_*M_, Lq_, P_, 2
          sampling_grid_l_ = (
              sampling_grids[:, :, :, level]
              .permute(0, 2, 1, 3, 4)
              .flatten(0, 1)
          )
          # N_*M_, D_, Lq_, P_
          sampling_value_l_ = torch.nn.functional.grid_sample(
              value_l_,
              sampling_grid_l_,
              mode="bilinear",
              padding_mode="zeros",
              align_corners=False,
          )
          sampling_value_list.append(sampling_value_l_)
  
      # (N_, Lq_, M_, L_, P_) -> (N_, M_, Lq_, L_, P_) -> (N_*M_, 1, Lq_, L_*P_)
      attention_weights = attention_weights.permute(
          0, 2, 1, 3, 4
      ).reshape(bs * n_head, 1, Len_q, n_levels * n_points)
      output = (
          (
              torch.stack(sampling_value_list, dim=-2).flatten(-2)
              * attention_weights
          )
          .sum(-1)
          .reshape(bs, n_head * c, Len_q)
      )
  
      return output.permute(0, 2, 1)
  
imports: from __future__ import annotations
  from typing import Optional
  import torch
  from torch import nn
  from kornia.contrib.models.common import MLP, ConvNormAct
  from kornia.core import Module, Tensor, concatenate
label_desc: Use image node_id 2 for calling a global function with example usage: _deformable_attention_kernel(value, value_spatial_shapes, sampling_locations, attention_weights) and returns: output