[Bug] [Relax][Torch] torch.logical_not on float tensor is lowered to R.logical_not with float dtype and fails LegalizeOps
- 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