pytorch / pytorch/pytorch

[inductor] mul_softmax_pattern generates neg(bool): compile-time NotImplementedError for softmax over a bool*tensor product (eager works)

Open Beginner friendly
#196,743 4 comments 0 reactions 0 assignees View on GitHub
bot-triaged module: inductor oncall: pt2 release triage triaged
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.