From 396ac876d521fd3a29ceccbd43266757a84b3311 Mon Sep 17 00:00:00 2001 From: mramanathan Date: Tue, 10 Dec 2024 11:28:27 +0000 Subject: [PATCH] Return dynamic inputs in DecomposeMultiDimSqueeze callback --- python/tvm/relay/op/contrib/forge/forge_passes.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/tvm/relay/op/contrib/forge/forge_passes.py b/python/tvm/relay/op/contrib/forge/forge_passes.py index 45f811aa5..a1f117c26 100644 --- a/python/tvm/relay/op/contrib/forge/forge_passes.py +++ b/python/tvm/relay/op/contrib/forge/forge_passes.py @@ -1883,6 +1883,8 @@ def __init__(self): def callback(self, pre, post, node_map): act = node_map[self.act][0] axis = post.attrs.axis + if any([isinstance(dim, tvm.tir.expr.Any) for dim in pre.checked_type.shape]): + return post input_shape = [int(dim) for dim in pre.args[0].checked_type.shape] adjusted_axes = [(ax - len(input_shape)) if ax >= 0 else ax for ax in axis] assert all(ax < 0 for ax in adjusted_axes), "Invalid squeeze dimension: all axes must be negative."