[Bug][Relax] RewriteDataflowReshape turns a resolved zero extent into a copied dimension

Start with src/relax/transform/rewrite_dataflow_reshape.cc, especially the RewriteDataflowReshape call to reshape(arg, res_ty->shape.value()). Read the reshape semantics that treat literal zero as a copy instruction, then reproduce with the supplied native Relax script and compare the ordinary build with RewriteDataflowReshape disabled. Done means the static-zero grid sample builds and returns shape (2,2,3,0) while the width-one control still passes.

Expected behavior

The native Relax graph below has image shape (2,2,3,2) and grid shape (2,3,0,2). Its grid sample output is the empty tensor (2,2,3,0). A reshape rewrite should preserve that zero extent.

Actual behavior

Both the empty graph and the neighboring width-one graph are well formed. The ordinary build of the empty graph fails with old shape [2,3,0,2], product 0 versus new shape [2,2,3,2], product 24. The stack reaches RewriteDataflowReshape. Disabling only that pass compiles and executes the identical empty module, returning the exact expected shape and dtype. The width-one control executes and matches independent NumPy arithmetic. There is no ONNX import or TVMScript parse/roundtrip in this reproducer.

Environment

• Ubuntu 24.04.1 LTS, x86_64; Python 3.12.3; NumPy 2.5.4; LLVM 18.1.3, CPU only.
• TVM reports 0.28.dev0; the imported Python snapshot is based on 27576ce2d1b2cf5212ad46d5710e594bdcbb5df0.
• The deployed compiler binary build commit was not retained. Its SHA-256 is 574a573fe994e5eda7352ddd8d685e64e45c7934bdeeaab57a5a272a77792ab1; this is not presented as a build of current main.
• Current main source was inspected at 2bf0532998d4dc7879cb342e873608b90c1d4497. I have not built or executed that revision.

Steps to reproduce

The complete script below was executed in two separate fresh processes. It contains the failing case and passing controls, uses only public synthetic data, and needs no model files or project-specific tools. Its assertions verify the reported behavior.

"""Native Relax only: static empty grid, nonempty control, and pass isolation."""
import numpy as np
import tvm
from tvm import relax


def run(width, disable=False):
    image = relax.Var('image', relax.TensorType((2, 2, 3, 2), 'float32'))
    grid = relax.Var('grid', relax.TensorType((2, 3, width, 2), 'float32'))
    bb = relax.BlockBuilder()
    with bb.function('main', [image, grid]):
        with bb.dataflow():
            transposed = bb.emit(relax.op.permute_dims(grid, [0, 3, 1, 2]))
            y = bb.emit_output(relax.op.image.grid_sample(
                image, transposed, align_corners=True, padding_mode='border'))
        bb.emit_func_output(y)
    module = bb.get()
    assert relax.analysis.check_well_formed(module, check_ty=True)
    values = np.arange(24, dtype='float32').reshape(2, 2, 3, 2)
    coordinates = np.zeros((2, 3, width, 2), dtype='float32')
    expected = np.broadcast_to(values[:, :, 1, :].mean(axis=-1)[:, :, None, None],
                               (2, 2, 3, width))
    try:
        with tvm.transform.PassContext(disabled_pass=['RewriteDataflowReshape'] if disable else []):
            executable = relax.build(module, target='llvm')
    except ValueError as error:
        assert width == 0 and not disable
        assert 'product T.int64(0)' in str(error) and 'product T.int64(24)' in str(error)
        print('width', width, 'build error:', str(error))
        return
    assert width != 0 or disable, 'Expected the reported static-zero build failure'
    output = relax.VirtualMachine(executable, tvm.cpu())['main'](
        tvm.runtime.tensor(values), tvm.runtime.tensor(coordinates)).numpy()
    assert output.shape == expected.shape and output.dtype == expected.dtype
    np.testing.assert_array_equal(output, expected)
    print('width', width, 'disabled pass', disable, 'exact output shape', output.shape)


run(0)
run(1)
run(0, disable=True)

Investigation and related work

RewriteDataflowReshape calls reshape(arg, res_ty->shape.value()) on an already resolved result shape. The semantic reshape constructor treats literal zero as a copy instruction, so this changes the target before type inference. #20255 concerns the Torch reshape caller; #20721 concerns ONNX shape handling and script printing, whereas this failure needs neither frontend nor script parsing. #19688 concerns 5D GridSample import, and #15089 concerns reflection-mode division by zero. An empty permutation alone also passed in my checks; the report is scoped to the demonstrated native graph, rather than every empty permutation.

Triage

• needs-triage
• type: bug

Codex assisted with report preparation; the included script was executed in the environment above.