Skip to content

[Bug][Relax][Frontend][Torch] Tensor.expand mishandles -1 when sizes has more dims than the input (right-alignment): valid models crash at import with IndexError #20229

Description

@siyiweigeHEW

Expected behavior

torch.Tensor.expand(*sizes) allows prepending new leading dimensions when len(sizes) > input.rank: the input shape is right-aligned under sizes, and -1 in sizes keeps the size of the corresponding input dimension. Both of these are valid and return the documented shape:

torch.zeros(2, 3).expand(4, -1, -1)  # -> (4, 2, 3)
torch.zeros(3).expand(2, -1)         # -> (2, 3)
torch.zeros(2, 3).expand(4, -1, 3)   # -> (4, 2, 3)

Actual behavior

tvm.relax.frontend.torch.from_exported_program crashes at model import for all of the above. The converter (_expand in python/tvm/relax/frontend/torch/base_fx_graph_translator.py:1752-1761) resolves each -1 by left-to-right index into the input shape:

for idx, i in enumerate(sizes):
    if isinstance(i, int) and i == -1:
        broadcast_shape.append(in_shape[idx])   # line 1758
    ...

It never right-aligns the input shape under sizes. When a -1 sits at a position past the input's rank (a newly prepended leading dim), in_shape[idx] is out of range and the whole model import fails:

(2,3).expand(4,-1,-1)  torch -> (4,2,3)   TVM IndexError: ShapeExpr index out of range
(3,).expand(2,-1)      torch -> (2,3)     TVM IndexError: ShapeExpr index out of range

When the wrong index is still in range, the -1 is mapped to the wrong input dimension, producing an invalid broadcast target:

(2,3).expand(4,-1,3)  torch -> (4,2,3)
# TVM maps -1 -> in_shape[1]=3 instead of in_shape[0]=2 -> broadcast_shape=[4,3,3]
TVM InternalError: broadcast_to expects the input tensor shape is broadcastable to the target ...

All six tested models are plain, runnable PyTorch modules.

Environment

  • OS: Linux
  • TVM: v0.24.dev0 (main branch, commit 262c6d2e0, built 2026-02-11)
  • Python: 3.11
  • torch: 2.10.0

Steps to reproduce

"""Repro: torch.Tensor.expand -1 / right-alignment mishandled by TVM relax torch frontend."""
import torch
import torch.nn as nn
from tvm.relax.frontend.torch import from_exported_program

class M(nn.Module):
    def forward(self, x):
        return x.expand(4, -1, -1)          # sizes longer than input rank

x = torch.randn(2, 3)
print("torch:", tuple(M()(x).shape))        # (4, 2, 3)

exp = torch.export.export(M().eval().cpu(), (x.cpu(),))
print(exp.graph)                            # %expand = aten.expand.default(%x, [4, -1, -1])
from_exported_program(exp)                  # TVM raises IndexError

Actual output:

torch: (4, 2, 3)
graph():
    %x : [num_users=1] = placeholder[target=x]
    %expand : [num_users=1] = call_function[target=torch.ops.aten.expand.default](args = (%x, [4, -1, -1]), kwargs = {})
    return (expand,)
Error converting operator expand, with inputs: [data, metadata["relax.expr.Constant"][0]]
TVM: IndexError: ShapeExpr index out of range
  File "tvm/relax/frontend/torch/base_fx_graph_translator.py", line 1758, in _expand
    broadcast_shape.append(in_shape[idx])

More cases (all pass in torch, all fail in TVM):

torch.zeros(3).expand(2, -1)             # torch (2,3)      TVM IndexError
torch.zeros(2, 3, 4).expand(5, -1, -1, -1)  # torch (5,2,3,4)  TVM IndexError
torch.zeros(1, 3).expand(4, -1, -1)      # torch (4,1,3)    TVM IndexError
torch.zeros(2, 3).expand(4, -1, 3)       # torch (4,2,3)    TVM InternalError (wrong -1 mapping)
torch.zeros(2, 1).expand(4, -1, -1)      # torch (4,2,1)    TVM IndexError

For comparison, the common cases where -1 stays within the input rank (no new leading dims) work correctly (max|diff|=0 vs torch): (1,3).expand(4,-1), (2,3,4).expand(2,-1,4), (2,3).expand(2,3).

Triage

  • needs-triage
  • bug
  • relax
  • frontend/torch

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    needs-triagePRs or issues that need to be investigated by maintainers to find the right assignees to address ittype: bug

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions