code: def _roi_align(
      input,
      rois,
      spatial_scale,
      pooled_height,
      pooled_width,
      sampling_ratio,
      aligned,
  ):
      orig_dtype = input.dtype
  
      input = maybe_cast(input)
      rois = maybe_cast(rois)
  
      _, _, height, width = input.size()
  
      ph = torch.arange(pooled_height, device=input.device)  # [PH]
      pw = torch.arange(pooled_width, device=input.device)  # [PW]
  
      # input: [N, C, H, W]
      # rois: [K, 5]
  
      roi_batch_ind = rois[:, 0].int()  # [K]
      offset = 0.5 if aligned else 0.0
      roi_start_w = rois[:, 1] * spatial_scale - offset  # [K]
      roi_start_h = rois[:, 2] * spatial_scale - offset  # [K]
      roi_end_w = rois[:, 3] * spatial_scale - offset  # [K]
      roi_end_h = rois[:, 4] * spatial_scale - offset  # [K]
  
      roi_width = roi_end_w - roi_start_w  # [K]
      roi_height = roi_end_h - roi_start_h  # [K]
      if not aligned:
          roi_width = torch.clamp(roi_width, min=1.0)  # [K]
          roi_height = torch.clamp(roi_height, min=1.0)  # [K]
  
      bin_size_h = roi_height / pooled_height  # [K]
      bin_size_w = roi_width / pooled_width  # [K]
  
      exact_sampling = sampling_ratio > 0
  
      roi_bin_grid_h = (
          sampling_ratio
          if exact_sampling
          else torch.ceil(roi_height / pooled_height)
      )  # scalar or [K]
      roi_bin_grid_w = (
          sampling_ratio
          if exact_sampling
          else torch.ceil(roi_width / pooled_width)
      )  # scalar or [K]
  
      """
      iy, ix = dims(2)
      """
  
      if exact_sampling:
          count = max(roi_bin_grid_h * roi_bin_grid_w, 1)  # scalar
          iy = torch.arange(roi_bin_grid_h, device=input.device)  # [IY]
          ix = torch.arange(roi_bin_grid_w, device=input.device)  # [IX]
          ymask = None
          xmask = None
      else:
          count = torch.clamp(
              roi_bin_grid_h * roi_bin_grid_w, min=1
          )  # [K]
          # When doing adaptive sampling, the number of samples we need to do
          # is data-dependent based on how big the ROIs are.  This is a bit
          # awkward because first-class dims can't actually handle this.
          # So instead, we inefficiently suppose that we needed to sample ALL
          # the points and mask out things that turned out to be unnecessary
          iy = torch.arange(height, device=input.device)  # [IY]
          ix = torch.arange(width, device=input.device)  # [IX]
          ymask = iy[None, :] < roi_bin_grid_h[:, None]  # [K, IY]
          xmask = ix[None, :] < roi_bin_grid_w[:, None]  # [K, IX]
  
      def from_K(t):
          return t[:, None, None]
  
      y = (
          from_K(roi_start_h)
          + ph[None, :, None] * from_K(bin_size_h)
          + (iy[None, None, :] + 0.5).to(input.dtype)
          * from_K(bin_size_h / roi_bin_grid_h)
      )  # [K, PH, IY]
      x = (
          from_K(roi_start_w)
          + pw[None, :, None] * from_K(bin_size_w)
          + (ix[None, None, :] + 0.5).to(input.dtype)
          * from_K(bin_size_w / roi_bin_grid_w)
      )  # [K, PW, IX]
      val = _bilinear_interpolate(
          input, roi_batch_ind, y, x, ymask, xmask
      )  # [K, C, PH, PW, IY, IX]
  
      # Mask out samples that weren't actually adaptively needed
      if not exact_sampling:
          val = torch.where(ymask[:, None, None, None, :, None], val, 0)
          val = torch.where(xmask[:, None, None, None, None, :], val, 0)
  
      output = val.sum((-1, -2))  # remove IY, IX ~> [K, C, PH, PW]
      if isinstance(count, torch.Tensor):
          output /= count[:, None, None, None]
      else:
          output /= count
  
      output = output.to(orig_dtype)
  
      return output
  
file_path: vision/torchvision/ops/roi_align.py
imports: from typing import List, Union
  import torch
  import torch._dynamo
  import torch.fx
  from torch import nn, Tensor
  from torch.jit.annotations import BroadcastingList2
  from torch.nn.modules.utils import _pair
  from torchvision.extension import _assert_has_ops, _has_ops
  from ..utils import _log_api_usage_once
  from ._utils import check_roi_boxes_shape, convert_boxes_to_roi_format
label_desc: Use image node_id 3 for calling a global function with example usage: _roi_align(input, rois, spatial_scale, pooled_height, pooled_width, sampling_ratio, aligned) and returns: output, t