Repository: pytorch
Issue: Some operators miss dtype check when using `torch.compile`
Labels: module: error checking,triaged,module: structured kernels,oncall: pt2,module: inductor

### 🐛 Describe the bug

As reported here (https://github.com/pytorch/pytorch/issues/144314#issuecomment-2574508557), I notice some operators missing dtype check when executed in the context of `torch.compile`. The specific symptom is as follows:
- Eager Mode: Raises `not implemented for [specific dtype]` error
- torch.compile Mode: Yields regular outputs (I guess implicit data type casting happens under `torch.compile`)

Some related issues: https://github.com/pytorch/pytorch/issues/144314, https://github.com/pytorch/pytorch/issues/144310, https://github.com/pytorch/pytorch/issues/144247.

Although this dtype-check-missing issue may not be severe, in case you are interested, I cherrypick a few operators where dtype checks are missing in the CPU and CUDA backends. Here's a breakdown:

| Operator Name |   CPU Backend Missing Check |   CUDA Backend Misses Check | Expected Behavior (Eager Behavior) |
| -------- | ------- | ------- | ------- | 
| torch.nn.functional.{log_softmax,softmax,logsigmoid} |    uint, int8, int16, int32, int64 |   uint, int8, int16, int32, int64 | Raise `not implemented for xxx` error |
| torch.nn.functional.{gelu,celu,hardsigmoid,hardswish}/torch.nextafter |   uint, bool, int8, int16, int32, int64 |   uint, bool, int8, int16, int32, int64 | Raise `not implemented for xxx` error |
| torch.nn.functional.prelu |   bool, int8, int16, int32, int64 |   uint, bool, int8, int16, int32, int64 | Raise `not implemented for xxx` error |
| torch.Tensor.mm |   uint, bool |   N/A | Raise `not implemented for xxx` error |
| torch.trace |   uint, bfloat16 , half, bool |   N/A | Raise `not implemented for xxx` error |
| torch.fmax |   complex32, complex64 |   N/A | Raise `not implemented for xxx` error |
| torch.xlogy/torch.nn.functional.mse_loss |   complex64, complex32 |   complex64, complex32 | Raise `not implemented for xxx` error |

Since these cases seem to share the same root cause, I am wondering if they can be fixed in a general way?

Below are detailed code that can reproduce the reported case for each operator.

<details>
<summary>log_softmax/softmax</summary>

```
import torch
from torch import nn
torch._dynamo.config.recompile_limit = 100
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()

    def forward(self, input,dim):
          return torch.nn.functional.log_softmax(input,dim)  # replace `log_softmax` with `softmax to reproduce the issue in softmax
f = MyModel()
cf = torch.compile(f)
input = torch.randn((2))
dim = -1

for device in ['cpu', 'cuda']:
    for dtype in [torch.uint16, torch.int8, torch.int16, torch.int32, torch.int64, torch.float16, torch.float32, torch.float64]:
        input = input.to(dtype).to(device)
        eager_pass, compile_pass = "passed", "passed"
        try:
            f(input, dim)
            eager_pass = "passed"
        except Exception as e:
            print(f"Eager Error: {e}")
            eager_pass = "failed"
        try:
            cf(input, dim)
            compile_pass = "passed"
        except Exception as e:
            compile_pass = "failed"
        if eager_pass != compile_pass:
            print(f"Inconsistent behavior on: {dtype}, {device}\n Eager: {eager_pass}\n Compile: {compile_pass}")
```
</details>


<details>
<summary>logsigmoid/gelu/celu/hardsigmoid/hardswish</summary>

```
import torch
from torch import nn
torch._dynamo.config.recompile_limit = 100

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()

    def forward(self, input):
          return torch.nn.functional.logsigmoid(input)  # change logsigmoid to gelu/celu/hardsigmoid/hardswish will reproduce related inconsistent behaviors

f = MyModel()
cf = torch.compile(f)
input = torch.randn((2))

for device in ['cpu', 'cuda']:
    for dtype in [torch.uint16, torch.bool, torch.int8, torch.int16, torch.int32, torch.int64, torch.float16, torch.float32, torch.float64]:
        input = input.to(dtype).to(device)
        eager_pass, compile_pass = "passed", "passed"
        try:
            f(input)
            eager_pass = "passed"
        except Exception as e:
            print(f"Eager Error: {e}")
            eager_pass = "failed"
        try:
            cf(input)
            compile_pass = "passed"
        except Exception as e:
            compile_pass = "failed"
        if eager_pass != compile_pass:
            print(f"Inconsistent behavior on: {dtype}, {device}\n Eager: {eager_pass}\n Compile: {compile_pass}")
```

</details>




<details>
<summary>prelu</summary>

```
import torch
from torch import nn
import numpy as np
torch._dynamo.config.recompile_limit = 100

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()

    def forward(self, input,weight):
          return torch.nn.functional.prelu(input,weight)

f = MyModel()
cf = torch.compile(f)
input = torch.tensor(np.random.randint(-10, 10, (1,1,1)))
weight = torch.tensor(np.random.randint(-10, 10, (1)))

for device in ['cpu', 'cuda']:
    for dtype in [torch.uint16, torch.bool, torch.int8, torch.int16, torch.int32, torch.int64, torch.float16, torch.float32, torch.float64]:
        input = input.to(dtype).to(device)
        weight = weight.to(dtype).to(device)
        eager_pass, compile_pass = "passed", "passed"
        try:
            f(input,weight)
            eager_pass = "passed"
        except Exception as e:
            print(f"Eager Error: {e}")
            eager_pass = "failed"
        try:
            cf(input,weight)
            compile_pass = "passed"
        except Exception as e:
            compile_pass = "failed"
        if eager_pass != compile_pass:
            print(f"Inconsistent behavior on: {dtype}, {device}\n Eager: {eager_pass}\n Compile: {compile_pass}")
```

</details>

<details>
<summary>torch.nextafter</summary>

```
import torch
from torch import nn
import numpy as np
torch._dynamo.config.recompile_limit = 100

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()

    def forward(self, input, other):
          return torch.nextafter(input, other)

f = MyModel()
cf = torch.compile(f)
input = torch.tensor(np.random.randint(-10, 10, ()), dtype=torch.int64)
other = torch.tensor(np.random.randint(-10, 10, ()), dtype=torch.int64)

for device in ['cpu', 'cuda']:
    for dtype in [torch.uint16, torch.bool, torch.int8, torch.int16, torch.int32, torch.int64, torch.float16, torch.float32, torch.float64]:
        input = input.to(dtype).to(device)
        other = other.to(dtype).to(device)
        eager_pass, compile_pass = "passed", "passed"
        try:
            f(input, other)
            eager_pass = "passed"
        except Exception as e:
            print(f"Eager Error: {e}")
            eager_pass = "failed"
        try:
            cf(input, other)
            compile_pass = "passed"
        except Exception as e:
            compile_pass = "failed"
        if eager_pass != compile_pass:
            print(f"Inconsistent behavior on: {dtype}, {device}\n Eager: {eager_pass}\n Compile: {compile_pass}")
```

</details>


<details>
<summary>torch.Tensor.mm</summary>

```
import torch
from torch import nn
torch._dynamo.config.recompile_limit = 100
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()

    def forward(self, input, mat2):
          return torch.Tensor.mm(input,mat2)
f = MyModel()
cf = torch.compile(f)
input = torch.randn(1, 1)
mat2 = torch.randn(1, 1)

for device in ['cpu', 'cuda']:
    for dtype in [torch.uint16, torch.bool, torch.int8, torch.int16, torch.int32, torch.int64, torch.float16, torch.float32, torch.float64]:
        input = input.to(dtype).to(device)
        mat2 = mat2.to(dtype).to(device)
        eager_pass, compile_pass = "passed", "passed"
        try:
            f(input, mat2)
            eager_pass = "passed"
        except Exception as e:
            print(f"Eager Error: {e}")
            eager_pass = "failed"
        try:
            cf(input, mat2)
            compile_pass = "passed"
        except Exception as e:
            compile_pass = "failed"
        if eager_pass != compile_pass:
            print(f"Inconsistent behavior on: {dtype}, {device}\n Eager: {eager_pass}\n Compile: {compile_pass}")
```

</details>

<details>
<summary>torch.trace</summary>

```
import torch
from torch import nn
torch._dynamo.config.recompile_limit = 100
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()

    def forward(self, input):
          return torch.trace(input)
f = MyModel()
cf = torch.compile(f)
input = torch.randn(0,1)

for device in ['cpu', 'cuda']:
    for dtype in [torch.uint16, torch.bool, torch.bfloat16, torch.half]:
        input = input.to(dtype).to(device)
        eager_pass, compile_pass = "passed", "passed"
        try:
            f(input)
            eager_pass = "passed"
        except Exception as e:
            print(f"Eager Error: {e}")
            eager_pass = "failed"
        try:
            cf(input)
            compile_pass = "passed"
        except Exception as e:
            compile_pass = "failed"
        if eager_pass != compile_pass:
            print(f"Inconsistent behavior on: {dtype}, {device}\n Eager: {eager_pass}\n Compile: {compile_pass}")
```

</details>


<details>
<summary>torch.fmax</summary>

```
import torch
from torch import nn
import numpy as np
torch._dynamo.config.recompile_limit = 100
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()

    def forward(self, input,other):
          return torch.fmax(input,other)
f = MyModel()
cf = torch.compile(f)
input = torch.tensor(np.random.randn(1,1,1), dtype=torch.complex128)
other = torch.tensor(np.random.randn(0), dtype=torch.double)

for device in ['cpu', 'cuda']:
    for dtype in [torch.uint16, torch.bool, torch.bfloat16, torch.half, torch.complex64, torch.complex128]:
        input = input.to(dtype).to(device)
        try:
            f(input, other)
            eager_pass = "passed"
        except Exception as e:
            print(f"Eager Error: {e}")
            eager_pass = "failed"
        try:
            cf(input, other)
            compile_pass = "passed"
        except Exception as e:
            compile_pass = "failed"
        if eager_pass != compile_pass:
            print(f"Inconsistent behavior on: {dtype}, {device}\n Eager: {eager_pass}\n Compile: {compile_pass}")
```

</details>



<details>
<summary>torch.xlogy/mse_loss</summary>

```
import torch
from torch import nn
import numpy as np
torch._dynamo.config.recompile_limit = 100
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()

    def forward(self, input,other):
          return torch.xlogy(input,other) # change torch.xlogy to torch.nn.functional.mse_loss can reproduce mse_loss's inconsistent behavior
f = MyModel()
cf = torch.compile(f)
input = torch.tensor(np.random.randn(1))
other = torch.tensor(np.random.randn(1,1))

for device in ['cpu', 'cuda']:
    for dtype in [torch.uint16, torch.bool, torch.bfloat16, torch.half, torch.complex64, torch.complex128]:
        input = input.to(dtype).to(device)
        other = other.to(dtype).to(device)
        try:
            f(input, other)
            eager_pass = "passed"
        except Exception as e:
            print(f"Eager Error: {e}")
            eager_pass = "failed"
        try:
            cf(input, other)
            compile_pass = "passed"
        except Exception as e:
            compile_pass = "failed"
        if eager_pass != compile_pass:
            print(f"Inconsistent behavior on: {dtype}, {device}\n Eager: {eager_pass}\n Compile: {compile_pass}")
```

</details>

To my best knowledge, I track other related issues here.

cc @malfet @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @aakhundov @BoyuanFeng
```[tasklist]
### Tasks
- [ ] https://github.com/pytorch/pytorch/issues/144314
- [ ] https://github.com/pytorch/pytorch/issues/144310
- [ ] https://github.com/pytorch/pytorch/issues/144247
- [ ] https://github.com/pytorch/pytorch/issues/143779
- [ ] https://github.com/pytorch/pytorch/issues/143801
- [ ] https://github.com/pytorch/pytorch/issues/143752
- [ ] https://github.com/pytorch/pytorch/issues/143729
```