apache / apache/tvm

[Bug] [Relax][Torch] torch.logical_not on float tensor is lowered to R.logical_not with float dtype and fails LegalizeOps

Open
#19,553 1 comment 0 reactions 0 assignees View on GitHub
needs-triage type: bug
Dominant language
Python
Stars
13.7k
Forks
4k
Avg merge
2d 1h
Merged PRs (30d)
112

Description

### Summary

A minimal PyTorch model using `torch.logical_not` on a float tensor can be exported to Relax, but fails during `tvm.compile(..., target="llvm")`.

PyTorch supports `torch.logical_not` on non-bool tensors. For example, a float input produces a bool output:

```python
x = torch.tensor([[0.0, 1.0]], dtype=torch.float32)
torch.logical_not(x)
# tensor([[ True, False]])
```

However, after `torch.export` and `from_exported_program`, TVM emits a Relax program where `R.logical_not` is applied directly to a `float32` tensor, and the output is also annotated as `float32`.

This later fails in `relax.transform.LegalizeOps`, because TOPI/TIR logical-not expects a boolean argument.

### Actual behavior

`tvm.compile` fails during Relax legalization:
```
tvm.error.InternalError: Check failed: (arg.dtype().is_bool()) is false:
Expected boolean argument for ! operator (logical NOT), but received x[ax0, ax1] of type float32
```
The relevant stack trace reaches:
```
relax.transform.LegalizeOps
-> tvm.topi.logical_not
-> tvm::logical_not
-> type_check_boolean_args
-> Expected boolean argument for ! operator (logical NOT)
```
### Environment

TVM: 0.23.0
LLVM: 17.0.6
Python: 3.10.16 (from stack paths)
NumPy: 2.2.6

### Steps to reproduce

```python
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import sys
import platform
import traceback

import torch
import tvm

class MyModel(torch.nn.Module):
def forward(self, x):
return torch.logical_not(x)

def main():
print("=" * 80)
print("Environment")
print("=" * 80)
print("python:", sys.version.replace("\n", " "))
print("platform:", platform.platform())
print("torch:", torch.__version__)
print("tvm:", getattr(tvm, "__version__", ""))
print("tvm path:", getattr(tvm, "__file__", ""))

model = MyModel().eval()
x = torch.tensor([[0.0, 1.0]], dtype=torch.float32)

with torch.no_grad():
eager = model(x)

print("=" * 80)
print("PyTorch eager")
print("=" * 80)
print("input:", x)
print("input dtype:", x.dtype, "shape:", tuple(x.shape))
print("eager output:", eager)
print("eager dtype:", eager.dtype, "shape:", tuple(eager.shape))

ep = torch.export.export(model, (x,))

from tvm.relax.frontend.torch import from_exported_program

ir_mod = from_exported_program(ep)

print("=" * 80)
print("Exported Relax IR")
print("=" * 80)
print(ir_mod.script(show_meta=True))

print("=" * 80)
print("tvm.compile with LLVM")
print("=" * 80)

ex = tvm.compile(
ir_mod,
target=tvm.target.Target("llvm"),
relax_pipeline="default",
tir_pipeline="default",
)

print("compile: OK")
print(ex)

if __name__ == "__main__":
try:
main()
except Exception:
print("compile: FAILED")
traceback.print_exc()
```

### Triage

* needs-triage
* bug

cc @junrushao @shingjan

Contributor guide

No contributing guide indexed for this repository

Research direction

Reproduce the failure with the provided PyTorch model and inspect the Relax IR produced by from_exported_program. Trace the R.logical_not call through relax.transform.LegalizeOps and the TOPI/TIR logical_not path, focusing on the float input and output annotations. Done means float tensors compile successfully and the result preserves PyTorch's boolean output semantics.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
compilers, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.