[inductor] mul_softmax_pattern generates neg(bool): compile-time NotImplementedError for softmax over a bool*tensor product (eager works)
- Dominant language
- Python
- Stars
- 103k
- Forks
- 29.5k
- PR merge metrics
- PR metrics pending
Description
### Summary
`torch.compile` (inductor) raises `NotImplementedError` at compile time for a model
that runs fine in eager, whenever a **bool** tensor is an operand of a multiply that
feeds a `softmax` (the mul→softmax numerical-stability rewrite). The inductor
`mul_softmax_pattern` joint-graph replacement builds `sign = torch.where(other >= 0,
one, -one)` with `one = torch.scalar_tensor(1, dtype=inp.dtype)`; when `inp` is bool,
`-one` negates a bool tensor and fails.
### PoC (eager OK; compiled fails on 2.14.0 and nightly 2.15.0.dev20260910)
```python
import torch
def f(m):
return torch.softmax(m * torch.tensor(2.0), dim=1) # m is bool
m = torch.randint(0, 2, (2, 4, 8), dtype=torch.bool)
f(m) # eager OK -> shape (2,4,8)
torch.compile(f)(m) # NotImplementedError: Negation, the `-` operator, on a bool tensor is not supported.
```
Controls:
- `m.float() * torch.tensor(2.0)` -> compiled OK (specific to bool dtype)
- `m * 2.0` (python scalar) -> compiled OK (scalar branch uses `1 if other>=0 else -1`, no neg)
### Details
`torch/_inductor/fx_passes/joint_graph.py`, `mul_softmax_pattern`'s `repl`:
```python
one = torch.scalar_tensor(1, dtype=inp.dtype, device=inp.device)
sign = torch.where(other >= 0, one, -one) # -one is neg on a bool tensor when inp is bool
inp = inp * sign
max_ = torch.amax(inp, dim=dim, keepdim=keepdim)
```
The pattern fires only when the scaling operand is broadcast-constant across the
softmax dim (`_other_is_broadcasted_in_dim`) and the `other` operand is a **tensor**
(the python-scalar path takes a different, safe branch). The fix is to build `one`
with a signed dtype (or guard the bool case / use logical ops) instead of inheriting
`inp.dtype`.
Traceback origin: `joint_graph.py ... in repl` -> fake-tensor `aten.neg` fallback ->
`NotImplementedError`.
### Versions
Reproduced on `2.14.0+cpu` and nightly `2.15.0.dev20260910+cpu` (git 682cb3eebadc). CPU, inductor.
### Dedup
- No existing issue about `mul_softmax_pattern` / softmax-mul bool negation (open or closed).
- Distinct from #195674 ("1 - mask": eager rejects, compiled evaluates). Opposite direction
(here eager accepts, compiled rejects) and different mechanism (joint-graph softmax rewrite
emitting neg-on-bool, not a subtraction decomposition).
cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo @eellison
Contributor guide
Research direction
Start by running the provided eager-versus-compiled reproduction, then inspect torch/_inductor/fx_passes/joint_graph.py and the mul_softmax_pattern repl function. Ensure the bool-tensor path no longer raises NotImplementedError during compilation, while the float and Python-scalar controls continue to compile successfully.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100