pytorch - 💡(How to fix) Fix [inductor] argmax over comparison mask is accepted and used for indexing while eager raises

🐛 Describe the bug

torch.argmax rejects bool input in eager mode. The same error is also preserved by torch.compile(..., backend="eager") and torch.compile(..., backend="aot_eager").

However, when the bool tensor is produced by a comparison expression such as scores > 0.0, torch.compile(..., backend="inductor", fullgraph=True) silently accepts torch.argmax on the bool mask. The resulting index is then consumed by downstream indexing, and the compiled graph returns a final numeric value.

This is an Inductor-only eager-parity issue. Compiled execution should raise the same RuntimeError as eager, rather than producing an index from an unsupported bool argmax.

Expected behavior: Inductor should raise RuntimeError: argmax(): does not support bool input.

Actual behavior: Inductor returns tensor(30.).

import torch
import traceback

def fn(scores, values):
    mask = scores > 0.0
    idx = torch.argmax(mask, dim=1)
    return values[idx].sum()

scores = torch.tensor(
    [[-1.0, 2.0, -3.0],
     [4.0, -5.0, 6.0]],
    dtype=torch.float32,
)

values = torch.tensor([10.0, 20.0, 30.0])

try:
    print("eager:", fn(scores, values))
except Exception as e:
    print("eager error:", type(e).__name__, str(e).splitlines()[0])

for backend in ["eager", "aot_eager", "inductor"]:
    try:
        cf = torch.compile(fn, backend=backend, fullgraph=True)
        print(backend, "compiled:", cf(scores, values))
    except Exception as e:
        print(backend, "compiled error:", type(e).__name__, str(e).splitlines()[0])
        traceback.print_exc(limit=4)

Output:

eager error: RuntimeError argmax(): does not support bool input
eager compiled error: RuntimeError argmax(): does not support bool input
aot_eager compiled error: RuntimeError argmax(): does not support bool input
inductor compiled: tensor(30.)

This may be related to existing bool argmin/argmax correctness work, but the issue here is specifically that eager rejects torch.argmax on bool input while Inductor silently accepts it and allows the unsupported result to affect downstream indexing.

Versions

PyTorch version: 2.10.0+cpu