Mojo function

roi_align_nhwc

roi_align_nhwc[dtype: DType, //, aligned: Bool, mode: StringSlice[StaticConstantOrigin] = "AVG"](output: TileTensor[dtype, LayoutType, origin, address_space=address_space, linear_idx_type=linear_idx_type, element_shape_types=element_shape_types], input: TileTensor[dtype, LayoutType, origin, address_space=address_space, linear_idx_type=linear_idx_type, element_shape_types=element_shape_types], rois: TileTensor[dtype, LayoutType, origin, address_space=address_space, linear_idx_type=linear_idx_type, element_shape_types=element_shape_types], output_height: Int, output_width: Int, in_spatial_scale: Scalar[dtype], in_sampling_ratio: Scalar[dtype])

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].

Parameters:

dtype (DType): Type of the input tensor.
aligned (Bool): If not true offset the ROIs by 0.5.
mode (StringSlice): The pooling mode "AVG" for average and "MAX" for max pooling.