Maxpool2d failure when lowering to stablehlo

import torch
import torch_mlir

class Module(torch.nn.Module):

    def __init__(self):
        super().__init__()

    def forward(self, val):
        return torch.ops.aten.max_pool2d_with_indices(val, [2, 2])

module = torch_mlir.compile(Module(), [torch.randn(128, 3, 64, 64)], output_type="stablehlo")
print(module.operation.get_asm())         

fails with

python exception: Failure while executing pass pipeline:
error: "aten::max_pool2d_with_indices"("maxpool2d_test.py":10:15): 'stablehlo.reduce_window' op inferred type(s) 'tensor<128x3x63x63xf32>', 'tensor<128x3x63x63xi64>' are incompatible with return type(s) of operation 'tensor<128x3x32x32xf32>', 'tensor<128x3x32x32xi64>'
error: "aten::max_pool2d_with_indices"("maxpool2d_test.py":10:15): 'stablehlo.reduce_window' op failed to infer returned types
note: "aten::max_pool2d_with_indices"("maxpool2d_test.py":10:15): see current operation:
%28:2 = "stablehlo.reduce_window"(%0, %26, %9, %27) ({
^bb0(%arg1: tensor<f32>, %arg2: tensor<i64>, %arg3: tensor<f32>, %arg4: tensor<i64>):
%31 = "stablehlo.compare"(%arg1, %arg3) {compare_type = #stablehlo<comparison_type FLOAT>, comparison_direction = #stablehlo<comparison_direction GE>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
%32 = "stablehlo.select"(%31, %arg1, %arg3) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
%33 = "stablehlo.compare"(%arg1, %arg3) {compare_type = #stablehlo<comparison_type FLOAT>, comparison_direction = #stablehlo<comparison_direction EQ>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
%34 = "stablehlo.minimum"(%arg2, %arg4) : (tensor<i64>, tensor<i64>) -> tensor<i64>
%35 = "stablehlo.select"(%31, %arg2, %arg4) : (tensor<i1>, tensor<i64>, tensor<i64>) -> tensor<i64>
%36 = "stablehlo.select"(%33, %34, %35) : (tensor<i1>, tensor<i64>, tensor<i64>) -> tensor<i64>
"stablehlo.return"(%32, %36) : (tensor<f32>, tensor<i64>) -> ()
}) {padding = dense<0> : tensor<4x2xi64>, window_dilations = dense<1> : tensor<4xi64>, window_dimensions = dense<[1, 1, 2, 2]> : tensor<4xi64>, window_strides = dense<1> : tensor<4xi64>} : (tensor<128x3x64x64xf32>, tensor<128x3x64x64xi64>, tensor<f32>, tensor<i64>) -> (tensor<128x3x32x32xf32>, tensor<128x3x32x32xi64>)

Start by running the provided max_pool2d_with_indices reproduction and inspect the lowering that creates stablehlo.reduce_window. Compare its inferred 63x63 result with the expected 32x32 output, then verify that compilation succeeds for the example and produces both returned tensors with the expected shapes.