text: from __future__ import absolute_import
  from __future__ import print_function
  from __future__ import division
  
  import warnings
  import math
  
  import torch
  from torch import nn
  import torch.nn.functional as F
  from torch.nn.init import xavier_uniform_, constant_
  
  from .ops.functions import MSDeformAttnFunction
  
  
  def _is_power_of_2(n):
      if (not isinstance(n, int)) or (n < 0):
          raise ValueError("invalid input for _is_power_of_2: {} (type: {})".format(n, type(n)))
      return (n & (n-1) == 0) and n != 0
  
  
  class MSDeformAttn(nn.Module):
      def __init__(self, d_model, n_levels, n_heads, n_points):
          """
          Multi-Scale Deformable Attention Module
          :param d_model      hidden dimension
          :param n_levels     number of feature levels
          :param n_heads      number of attention heads
          :param n_points     number of sampling points per attention head per feature level
          """
          super().__init__()
          if d_model % n_heads != 0:
              raise ValueError('d_model must be divisible by n_heads, but got {} and {}'.format(d_model, n_heads))
          _d_per_head = d_model // n_heads
          # you'd better set _d_per_head to a power of 2 which is more efficient in our CUDA implementation
          if not _is_power_of_2(_d_per_head):
              warnings.warn("You'd better set d_model in MSDeformAttn to make the dimension of each attention head a power of 2 "
                            "which is more efficient in our CUDA implementation.")
  
          self.im2col_step = 64
  
          self.d_model = d_model
          self.n_levels = n_levels
          self.n_heads = n_heads
          self.n_points = n_points
  
          self.sampling_offsets = nn.Linear(d_model, n_heads * n_levels * n_points * 2)
          self.attention_weights = nn.Linear(d_model, n_heads * n_levels * n_points)
          self.value_proj = nn.Linear(d_model, d_model)
          self.output_proj = nn.Linear(d_model, d_model)
  
          self._reset_parameters()
  
      def _reset_parameters(self):
          constant_(self.sampling_offsets.weight.data, 0.)
          thetas = torch.arange(self.n_heads, dtype=torch.float32) * (2.0 * math.pi / self.n_heads)
          grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)
          grid_init = (grid_init / grid_init.abs().max(-1, keepdim=True)[0]).view(self.n_heads, 1, 1, 2).repeat(1, self.n_levels, self.n_points, 1)
          for i in range(self.n_points):
              grid_init[:, :, i, :] *= i + 1
          with torch.no_grad():
              self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1))
          constant_(self.attention_weights.weight.data, 0.)
          constant_(self.attention_weights.bias.data, 0.)
          xavier_uniform_(self.value_proj.weight.data)
          constant_(self.value_proj.bias.data, 0.)
          xavier_uniform_(self.output_proj.weight.data)
          constant_(self.output_proj.bias.data, 0.)
  
      def forward(self, query, reference_points, input_flatten, input_spatial_shapes, input_level_start_index):
          """
          :param query                       (N, Length_{query}, C)
          :param reference_points            (N, Length_{query}, n_levels, 2), range in [0, 1], top-left (0,0), bottom-right (1, 1)
          :param input_flatten               (N, \sum_{l=0}^{L-1} H_l \cdot W_l, C)
          :param input_spatial_shapes        (n_levels, 2), [(H_0, W_0), (H_1, W_1), ..., (H_{L-1}, W_{L-1})]
          :param input_level_start_index     (n_levels, ), [0, H_0*W_0, H_0*W_0+H_1*W_1, H_0*W_0+H_1*W_1+H_2*W_2, ..., H_0*W_0+H_1*W_1+...+H_{L-1}*W_{L-1}]
          
          :return output                     (N, Length_{query}, C)
          """
          N, Len_q, _ = query.shape
          N, Len_in, _ = input_flatten.shape
          assert (input_spatial_shapes[:, 0] * input_spatial_shapes[:, 1]).sum() == Len_in
          
          value = self.value_proj(input_flatten)
          
          value = value.view(N, Len_in, self.n_heads, self.d_model // self.n_heads)
          sampling_offsets = self.sampling_offsets(query).view(N, Len_q, self.n_heads, self.n_levels, self.n_points, 2)
          attention_weights = self.attention_weights(query).view(N, Len_q, self.n_heads, self.n_levels * self.n_points)
          attention_weights = F.softmax(attention_weights, -1).view(N, Len_q, self.n_heads, self.n_levels, self.n_points)
          
          # N, Len_q, n_heads, n_levels, n_points, 2
          if reference_points.shape[-1] == 2:
              offset_normalizer = torch.stack([input_spatial_shapes[..., 1], input_spatial_shapes[..., 0]], -1)
              
              sampling_locations = reference_points[:, :, None, :, None, :] \
                                   + sampling_offsets / offset_normalizer[None, None, None, :, None, :]
          else:
              raise ValueError(
                  'Last dim of reference_points must be 2 , but get {} instead.'.format(reference_points.shape[-1]))
          
          output = MSDeformAttnFunction.apply(
              value, input_spatial_shapes, input_level_start_index, sampling_locations, attention_weights, self.im2col_step)
          output = self.output_proj(output)
          return output
  
  
  
  