diff --git a/python/tvm/topi/image/resize.py b/python/tvm/topi/image/resize.py index 8d0a984982e1..69f3fa85eeb8 100644 --- a/python/tvm/topi/image/resize.py +++ b/python/tvm/topi/image/resize.py @@ -152,7 +152,8 @@ def get_inx( if coordinate_transformation_mode == "half_pixel": in_x = (x + 0.5) * scale_x - 0.5 elif coordinate_transformation_mode == "align_corners": - in_x = (image_width - 1).astype("float") / (target_width - 1) * x + # A singleton output has only x=0, which maps to the first input sample. + in_x = (image_width - 1).astype("float") / te.max(target_width - 1, 1) * x elif coordinate_transformation_mode == "asymmetric": if use_int_div: in_x = te.div(x, te.div(target_width, image_width)) diff --git a/tests/python/relax/test_frontend_from_exported_program.py b/tests/python/relax/test_frontend_from_exported_program.py index ef266b2be05a..233b44e4076d 100644 --- a/tests/python/relax/test_frontend_from_exported_program.py +++ b/tests/python/relax/test_frontend_from_exported_program.py @@ -4200,6 +4200,34 @@ def main(data: R.Tensor((3, 1, 3), dtype="float32")) -> R.Tuple( verify_model(Unbind2(), single_dim_args, {}, expected3) +@pytest.mark.skipif(not device_enabled("llvm"), reason="LLVM is required") +@pytest.mark.parametrize( + "input_hw, size, scale_factor, align_corners", + [ + ((4, 4), (1, 6), None, True), + ((4, 4), (6, 1), None, True), + ((4, 4), (1, 1), None, True), + ((1, 4), None, 1.5, True), + ((4, 4), (2, 3), None, True), + ((4, 4), (1, 6), None, False), + ], +) +def test_interpolate_bilinear_singleton_axis(input_hw, size, scale_factor, align_corners): + class Interpolate(Module): + def forward(self, x): + return torch.nn.functional.interpolate( + x, + size=size, + scale_factor=scale_factor, + mode="bilinear", + align_corners=align_corners, + ) + + height, width = input_hw + x = torch.arange(height * width, dtype=torch.float32).reshape(1, 1, height, width) + verify_model_numerically(Interpolate(), (x,), rtol=1e-6, atol=1e-6) + + def test_interpolate(): class InterpolateBilinear(Module): def forward(self, input): diff --git a/tests/python/te/test_te_create_function.py b/tests/python/te/test_te_create_function.py index 9c1fa4ad5129..55592337620d 100644 --- a/tests/python/te/test_te_create_function.py +++ b/tests/python/te/test_te_create_function.py @@ -808,6 +808,23 @@ def test_resize2d_symbolic(): _check_workload(te_resize2d_symbolic, tir_resize2d_symbolic, index_dtype_override="int64") +@pytest.mark.skipif(not tvm_tirx.testing.device_enabled("llvm"), reason="LLVM is required") +def test_resize1d_align_corners_symbolic_width(): + width = tirx.Var("width", "int64") + data = te.placeholder((1, 1, 5), "float32", name="data") + resized = topi.image.resize1d( + data, (0.0, 0.0), (width,), coordinate_transformation_mode="align_corners" + ) + resize = tvm.compile(te.create_function([data, resized]), target="llvm") + + values = tvm_tirx.runtime.tensor(np.arange(5, dtype="float32").reshape(1, 1, 5)) + for output_width in (1, 3): + actual = tvm_tirx.runtime.empty((1, 1, output_width), "float32") + resize(values, actual) + expected = np.linspace(0, 4, output_width, dtype="float32").reshape(1, 1, output_width) + np.testing.assert_array_equal(actual.numpy(), expected, strict=True) + + def test_extern_with_explicit_buffer_access(): def te_extern(): A = te.placeholder((128, 128), name="A")