[Bug][Relax][TOPI] GatherElements rejects valid empty indices during lowering

Start with the gather_elements legalizer in python/tvm/relax/transform/legalize_ops/manipulate.py, which calls TOPI gather, then the indices-axis check near line 1529 of include/tvm/topi/transform.h. Run the attached repro.py to confirm the 'indices_dim_i >= 1' failure for the (0,3) case. Done means the empty case builds and returns a float (0,3) array that matches NumPy and ONNX Runtime, while the (1,3) control still passes.

Summary

GatherElements(axis=0) with data shape [2,3] and indices shape [0,3] imports successfully but fails during default LLVM lowering with indices_dim_i >= 1 (0 vs. 1). Replacing the empty index dimension with 1 and filling the index tensor with zero executes correctly.

Environment

Linux x86_64 CPU/LLVM; Python 3.12.3; TVM 0.28.dev0; LLVM 18.1.3; ONNX 1.22.0; NumPy 2.5.4; ONNX Runtime 1.29.0 (CPUExecutionProvider).

The installed TVM binary's exact build source revision is UNKNOWN. Source links are a separate static inspection of main at 70ec7d5bf20326ef301ee3f06eb6aed51760331d; that revision was not built or executed for this report.

Loaded binary SHA256 identities:

• libtvm_ffi.so: 12072d0ccd9d5cb2c6dac80d42f4217ab50dbf795db03197000614a7a770a816
• libtvm_compiler.so: 574a573fe994e5eda7352ddd8d685e64e45c7934bdeeaab57a5a272a77792ab1
• libtvm_runtime.so: 4175916fc0ce6266c912629d2255094b94087425c6ab53ffd5e004d4c292c828
• libtvm_runtime_extra.so: 79c93351fca12dd7851964890cea7da65dd2b2e0e64e5c28bfed56c878b70817

Expected and observed behavior

Return an empty FLOAT tensor of shape (0,3). TVM rejects it before execution. The (1,3) indices control returns [[0,1,2]].

GatherElements-13 specifies an output with the indices shape. An empty indices tensor contains no out-of-bounds index. NumPy take_along_axis and unoptimized ORT CPU both give shape [0,3]. ONNX ReferenceEvaluator 1.22.0 represents this particular empty result as (0,), so the script does not use that implementation as the shape oracle.

Standalone reproduction

Save this as repro.py in an LLVM-enabled TVM environment, then run:

OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 python repro.py

from pathlib import Path
import hashlib
import numpy as np
import onnx
import onnxruntime as ort
import tvm
from onnx import TensorProto as T, helper as h, numpy_helper as nh
from onnx.reference import ReferenceEvaluator
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx


def make(nodes, inputs, outputs, initializers=(), opset=18):
    model = h.make_model(h.make_graph(nodes, "repro", inputs, outputs, list(initializers)),
                         opset_imports=[h.make_opsetid("", opset)], ir_version=10)
    onnx.checker.check_model(model, full_check=True)
    return model


def tvm_run(model, feeds):
    mod = from_onnx(model, keep_params_in_input=False)
    names = [value.name for value in model.graph.input]
    executable = relax.build(mod, target="llvm", relax_pipeline="default")
    result = relax.VirtualMachine(executable, tvm.cpu())["main"](
        *[tvm.runtime.tensor(feeds[name]) for name in names])
    return [result.numpy()] if hasattr(result, "numpy") else [v.numpy() for v in result]


def save(label, model, feeds):
    folder = Path("artifacts") / label
    folder.mkdir(parents=True, exist_ok=True)
    onnx.checker.check_model(model, full_check=True)
    onnx.save(model, folder / "model.onnx")
    np.savez(folder / "inputs.npz", **feeds)
    print(label, "model_sha256", hashlib.sha256((folder / "model.onnx").read_bytes()).hexdigest())
    print(label, "inputs_sha256", hashlib.sha256((folder / "inputs.npz").read_bytes()).hexdigest())


def reference(model, feeds, expected, check_ort=True):
    outputs = ReferenceEvaluator(model).run(None, feeds)
    for a, b in zip(outputs, expected, strict=True):
        assert a.shape == b.shape and a.dtype == b.dtype, (a.shape, a.dtype, b.shape, b.dtype)
        np.testing.assert_allclose(a, b, rtol=1e-5, atol=1e-6)
    if check_ort:
        opts = ort.SessionOptions()
        opts.intra_op_num_threads = opts.inter_op_num_threads = 1
        opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
        outputs = ort.InferenceSession(model.SerializeToString(), opts,
            providers=["CPUExecutionProvider"]).run(None, feeds)
        for a, b in zip(outputs, expected, strict=True):
            assert a.shape == b.shape and a.dtype == b.dtype
            np.testing.assert_allclose(a, b, rtol=1e-5, atol=1e-6)


def error_line(error):
    return type(error).__name__ + ": " + str(error).splitlines()[-1]


def same(actual, expected):
    assert actual.shape == expected.shape and actual.dtype == expected.dtype
    np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-6)


data = np.arange(6, dtype="float32").reshape(2, 3)
for size in (0, 1):
    indices = np.zeros((size, 3), dtype="int64")
    feeds = {"data": data, "indices": indices}
    expected = np.take_along_axis(data, indices, axis=0)
    model = make([h.make_node("GatherElements", ["data", "indices"], ["y"], axis=0)],
        [h.make_tensor_value_info("data", T.FLOAT, [2, 3]),
         h.make_tensor_value_info("indices", T.INT64, [size, 3])],
        [h.make_tensor_value_info("y", T.FLOAT, [size, 3])], opset=13)
    save(str(size), model, feeds)
    opts = ort.SessionOptions()
    opts.intra_op_num_threads = opts.inter_op_num_threads = 1
    opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
    actual_ort = ort.InferenceSession(model.SerializeToString(), opts,
        providers=["CPUExecutionProvider"]).run(None, feeds)[0]
    same(actual_ort, expected)
    # Pass concrete shapes to isolate lowering from the separate zero-shape importer bug.
    mod = from_onnx(model, shape_dict={k: v.shape for k, v in feeds.items()}, keep_params_in_input=False)
    try:
        executable = relax.build(mod, target="llvm", relax_pipeline="default")
        result = relax.VirtualMachine(executable, tvm.cpu())["main"](
            tvm.runtime.tensor(data), tvm.runtime.tensor(indices)).numpy()
    except Exception as error:
        assert size == 0 and "indices_dim_i >= 1" in str(error), str(error)
        print("empty indices", error_line(error), "expected shape", expected.shape)
    else:
        assert size == 1, "Empty-index lowering failure did not reproduce"
        same(result, expected)
        print("nonempty indices control passed", result.tolist())

Recorded output

empty indices InternalError: Check failed: indices_dim_i >= 1 (0 vs. 1) :  expected shape (0, 3)
nonempty indices control passed [[0.0, 1.0, 2.0]]

Causal isolation and related work

The gather_elements legalizer calls TOPI gather. Its static indices-axis check requires at least one index, which is stronger than this ONNX contract. The repro passes explicit concrete shapes to isolate that condition from the separate zero-extent importer issue.

#20217 / #20219 address negative-index normalization, and #20297 addresses out-of-bounds accesses. This case has no index values or memory accesses and reaches a different, nonempty-shape precondition.

This self-contained script is a newly constructed minimal diagnostic from the same graph-testing investigation, not a byte-identical replay of an archived model. It needs only TVM, ONNX, NumPy and ONNX Runtime, downloads no files, and writes its generated models and inputs to artifacts/. The assertions check both the observed failure/wrong result and the normal controls; exit code zero means that diagnostic reproduced. No global novelty or current-main execution claim is made.

I am happy to help add an in-tree regression test or check a proposed correction.