Mojo function

roi_align_nhwc

roi_align_nhwc[type: DType, input_shape: DimList, roi_shape: DimList, //, aligned: Bool, mode: StringSlice[StaticConstantOrigin] = __init__[__mlir_type.!kgen.string]("AVG")](output: NDBuffer[type, 4, origin, shape, strides], input: NDBuffer[type, 4, origin, input_shape, strides], rois: NDBuffer[type, 2, origin, roi_shape, strides], output_height: Int, output_width: Int, in_spatial_scale: SIMD[dtype, 1], in_sampling_ratio: SIMD[dtype, 1])

Compute ROIAlign a batch of rois of shape [M, 5] where the first dim is the batch index, followed by region box coordinates (y0, x0) (y1, x1). For inputs of NHWC format. The output shape is [M, output_height, output_width, C].

Paramerers: type: Type of the input tensor. input_shape: Shape of the input tensor. roi_shape: Shape of regions of interests (ROI). aligned: If not true offset the ROIs by 0.5. mode: The pooling mode "AVG" for average and "MAX" for max pooling.

Args:

output (NDBuffer[type, 4, origin, shape, strides]): Pre-allocated output tensor.
input (NDBuffer[type, 4, origin, input_shape, strides]): Batched images to the roi_align with NHWC format.
rois (NDBuffer[type, 2, origin, roi_shape, strides]): Batched ROIs box coordinates.
output_height (Int): Pooled output height.
output_width (Int): Pooled output width.
in_spatial_scale (SIMD[dtype, 1]): Scale factor to remap the roi_align coordinates to the input coordinates.
in_sampling_ratio (SIMD[dtype, 1]): Number of sampling points in the interpolation grid used to compute the output value of each pooled bin.