From 4113daa1fb5a7dfe54e1101d3960d2fe22c9517b Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 4 Nov 2026 23:03:25 +0100 Subject: [PATCH] [Fix][Relax] Preserve global median keepdims shape --- .../transform/legalize_ops/statistical.py | 4 ++ ..._transform_legalize_ops_median_keepdims.py | 61 +++++++++++++++++++ 2 files changed, 65 insertions(+) create mode 100644 tests/python/relax/test_transform_legalize_ops_median_keepdims.py diff --git a/python/tvm/relax/transform/legalize_ops/statistical.py b/python/tvm/relax/transform/legalize_ops/statistical.py index 7821ec8f1952..2d96c558a55e 100644 --- a/python/tvm/relax/transform/legalize_ops/statistical.py +++ b/python/tvm/relax/transform/legalize_ops/statistical.py @@ -104,6 +104,7 @@ def _te_median( ) -> te.Tensor | tuple[te.Tensor, te.Tensor]: # currently only supports one axis or no axis ~ same pytorch # todo: support multiple axis ~ same numpy + input_ndim = len(x.shape) shape_prod = _compute_shape_prod(x, axis) mid_index = (shape_prod - 1) // 2 @@ -123,6 +124,9 @@ def _te_median( median_idx = topi.gather(index_sorted, axis=ax, indices=indices) if axis is None or len(axis) == 0: + if axis is None and keepdims: + # Flattening for the sort must not change the rank of the result. + return topi.reshape(median_val, [1] * input_ndim) return median_val if keepdims else topi.squeeze(median_val, axis=axis) val = median_val diff --git a/tests/python/relax/test_transform_legalize_ops_median_keepdims.py b/tests/python/relax/test_transform_legalize_ops_median_keepdims.py new file mode 100644 index 000000000000..bbc00e52a865 --- /dev/null +++ b/tests/python/relax/test_transform_legalize_ops_median_keepdims.py @@ -0,0 +1,61 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Regression tests for median legalization.""" + +import numpy as np +import pytest + +import tvm +import tvm_tirx +from tvm import relax + + +def make_median(shape, keepdims): + bb = relax.BlockBuilder() + x = relax.Var("x", relax.TensorType(shape, "float32")) + with bb.function("main", [x]): + with bb.dataflow(): + result = bb.emit(relax.op.median(x, axis=None, keepdims=keepdims)) + result = bb.emit_output(result) + bb.emit_func_output(result) + return bb.get() + + +@pytest.mark.parametrize("shape", [(), (5,), (2, 3), (2, 3, 4)]) +@pytest.mark.parametrize("keepdims", [False, True]) +def test_median_all_axes_preserves_inferred_return_type(shape, keepdims): + original = make_median(shape, keepdims) + lowered = relax.transform.LegalizeOps()(original) + tvm_tirx.ir.assert_structural_equal(original["main"].ret_ty, lowered["main"].ret_ty) + + +@pytest.mark.skipif(not tvm_tirx.runtime.enabled("llvm"), reason="LLVM not enabled") +@pytest.mark.parametrize( + "shape,keepdims", [((), True), ((2, 3), True), ((2, 3, 4), True), ((2, 3), False)] +) +def test_median_all_axes_runtime_shape_and_value(shape, keepdims): + data = np.arange(int(np.prod(shape)), dtype="float32")[::-1].reshape(shape) + target = tvm_tirx.target.Target({"kind": "llvm", "mcpu": "generic"}) + with target: + exe = tvm.compile(make_median(shape, keepdims), target=target) + dev = tvm_tirx.cpu(0) + vm = relax.VirtualMachine(exe, dev) + result = vm["main"](tvm_tirx.runtime.tensor(data, dev)).numpy() + expected = np.asarray(np.sort(data.ravel())[(data.size - 1) // 2], dtype="float32") + expected = expected.reshape([1] * len(shape) if keepdims else ()) + assert result.shape == expected.shape + np.testing.assert_array_equal(result, expected)