Skip to content

[Bug] InternalError "Cannot reduce a 0 dim Tensor" when compiling a reduction over a 0-D (scalar) tensor #19676

Description

@raindrops-0199

Expected behavior

Compiling a model that reduces a 0-dimensional (scalar) tensor — e.g. torch.max(x.sum()) — should succeed. PyTorch and NumPy both define a reduction over a scalar as the identity (it returns the scalar unchanged):

>>> torch.max(torch.randn(4).sum())   # input is 0-D
tensor(-1.5686)

TVM should either lower a reduction over a 0-D tensor as the identity, or raise a clear frontend error — not hit an internal assertion at compile time.

Actual behavior

The model is accepted by torch.export and by relax.frontend.torch.from_exported_program, but tvm.compile aborts with an internal check failure during legalization:

InternalError: Check failed: ndim != 0 (0 vs. 0) : Cannot reduce a 0 dim Tensor
  File ".../tvm/include/tvm/topi/reduction.h", line 187, in tvm::topi::CommReduce(
        const tvm::te::Tensor&, const Optional<Array<Integer>>&, FReduce, bool, bool)
    TVM_FFI_ICHECK_NE(ndim, 0) << "Cannot reduce a 0 dim Tensor";

torch.max(x) (no dim) is converted to a reduction over all axes; with a 0-D input, topi::CommReduce asserts ndim != 0 instead of treating it as identity.

Environment

  • TVM: 0.25.dev0
  • OS: Ubuntu 22.04 (Linux 6.8, x86_64)
  • Python: 3.12
  • PyTorch: 2.12.0
  • target: llvm

Steps to reproduce

import torch
import torch.nn as nn
from torch.export import export
import tvm
from tvm import relax
from tvm.relax.frontend.torch import from_exported_program


class M(nn.Module):
    def forward(self, x):
        return torch.max(x.sum())   # x.sum() -> 0-D scalar; torch.max reduces it again


m = M().eval()
args = (torch.randn(4),)
ep = export(m, args)
mod = from_exported_program(ep, keep_params_as_input=True, unwrap_unit_return_tuple=True)
mod, _ = relax.frontend.detach_params(mod)
tvm.compile(mod, target=tvm.target.Target("llvm"))   

Triage

  • needs-triage
  • relax

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
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions