[Relax][PyTorch] Add decomposed operator support for interpolate

python/tvm/relax/frontend/torch/base_fx_graph_translator.py

@@ -1693,20 +1693,42 @@ def _index_tensor(self, node: fx.Node) -> relax.Var:
             axis, index_tensor = non_none_indices[0]
             return self.block_builder.emit(relax.op.take(data, index_tensor, axis=axis))
 
-        # General case: multiple non-None indices require advanced indexing
+        # Check if all indices can be squeezed to 1D for sequential take
+        def is_squeezable(idx):
+            if idx.struct_info.ndim == 1:
+                return True
+            if idx.struct_info.ndim == 2:
+                shape = idx.struct_info.shape
+                for d in shape:
+                    if isinstance(d, int) and d == 1:
+                        return True
+                    # Check for tir.IntImm
+                    if hasattr(d, "value") and d.value == 1:
+                        return True
+            return False
+
+        all_squeezable = all(is_squeezable(idx) for _, idx in non_none_indices)
+        if all_squeezable:
+            result = data
+            for axis, idx in reversed(non_none_indices):
+                if idx.struct_info.ndim > 1:
+                    idx = self.block_builder.emit(relax.op.squeeze(idx))
+                result = self.block_builder.emit(relax.op.take(result, idx, axis=axis))
+            return result
+
+        # General case: replace None with arange, reshaped for broadcasting
+        max_ndim = max((idx.struct_info.ndim for _, idx in non_none_indices), default=1)
         processed_indices = []
         data_shape = self.shape_of(data)
 
         for i, idx in enumerate(indices):
             if idx is None:
-                dim_size = data_shape[i]
                 arange_idx = self.block_builder.emit(
-                    relax.op.arange(
-                        start=relax.PrimValue(0),
-                        end=dim_size,
-                        step=relax.PrimValue(1),
-                        dtype="int64",
-                    )
+                    relax.op.arange(relax.PrimValue(0), data_shape[i], relax.PrimValue(1), "int64")
+                )
+                # Reshape to [dim_size, 1, 1, ...] for broadcasting
+                arange_idx = self.block_builder.emit(
+                    relax.op.reshape(arange_idx, [data_shape[i]] + [1] * (max_ndim - 1))
                 )
                 processed_indices.append(arange_idx)
             else: