(tvm) branch main updated: [TOPI] Fix index handling in expand_like operator for axis expansion (#18006)

[mshr]

This is an automated email from the ASF dual-hosted git repository.

mshr pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 55b8980a3e [TOPI] Fix index handling in expand_like operator for axis 
expansion (#18006)
55b8980a3e is described below

commit 55b8980a3e7e60ce4a0ac48e60c61a1ac312ff14
Author: Qingchao Shen <qingchaos...@outlook.com>
AuthorDate: Sun Jun 1 19:22:09 2025 +0800

    [TOPI] Fix index handling in expand_like operator for axis expansion 
(#18006)
---
 python/tvm/topi/transform.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/python/tvm/topi/transform.py b/python/tvm/topi/transform.py
index 1ef6523059..951944e618 100644
--- a/python/tvm/topi/transform.py
+++ b/python/tvm/topi/transform.py
@@ -96,7 +96,8 @@ def expand_like(a, shape_like, axis):
         axis_index = 0
         for i in range(0, len(idxs)):
             if i not in real_axis:
-                indices.append(idxs[i])
+                dim = tvm.tir.if_then_else(a.shape[len(indices)] != 1, 
idxs[i], 0)
+                indices.append(dim)
                 axis_index += 1
         return a(*indices)