Repository: pytorch
Issue: linear mm weight and bias dtypes mismatch bypasses
Labels: oncall: jit

### 🐛 Describe the bug

Should linear weight and bias must be of the same dtypes? 

For linear, I see for {2,3}d rank tensor, weight and bias dtype mismatch is checked in [addmm](https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/LinearAlgebra.cpp#L155-L169). However, for higher rank tensor, weight and bias dtype mismatch is not checked in [mm](https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/LinearAlgebra.cpp#L175). I wonder if 
- this is unexpected, and we should add a similar check for dtype mismatch in addmm to mm
- this is expected (weight and bias can have different dtypes)


=========================================
Original Context:
Similar to conv-bn folding #76578, lin-bn folding #86706 with autocast freeze on gpu casts the [inputs to linear to half](https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/passes/autocast.cpp#L363), while inputs to batchnorm are not casted, remaining as float. So if the optional linear bias is set to false, bias will remain as float while weight is casted to half during autocast jit freeze, causing inconsistency in linear and bias dtype (half, float respectively). 

However, this inconsistency is not checked in linear admm (is checked in mm). I wonder if
- this is unexpected, and linear weight and bias must be of the same dtypes
for {2,3}d tensors, dtype mistmatch is checked in [addmm](https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/LinearAlgebra.cpp#L155-L169); for higher rank tensors, dtype mismatch is not checked in [mm](https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/LinearAlgebra.cpp#L175). If this is unexpected, we can add a similar check as in addmm to mm. 
- this is expected, bias can be kept in fp32 for accuracy reason.

```python
import torch
import torch.nn as nn

class LinearBN(torch.nn.Module):
    def __init__(self, in_features, out_features):
        super(LinearBN, self).__init__()
        self.linear = nn.Linear(in_features, out_features, bias=False)
        self.bn = nn.BatchNorm2d(out_features)
    def forward(self, x):
        x = self.linear(x)
        return self.bn(x)
net = LinearBN(32, 32).cuda().eval()  
print("weight: ", net.linear.weight.dtype) #torch.float32
img = torch.rand((1, 32, 32, 32)).cuda()
with torch.cuda.amp.autocast(True):
    net = torch.jit.script(net)
    net = torch.jit.freeze(net)
lin_node = net.graph.findNode("aten::linear", True)
print("weight: ", lin_node.namedInput("weight").type().dtype()) #torch.float16
print("bias: ", lin_node.namedInput("bias").type().dtype()) #torch.float32
net(img) #passes without catching weight, bias dtype inconsistency, expected?
```

For 2/3d tensors, the dtypes mismatch is catched in [addmm](https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/LinearAlgebra.cpp#L155-L169). 
```python
class LinearBN(torch.nn.Module):
    def __init__(self, in_features, out_features):
        super(LinearBN, self).__init__()
        self.linear = nn.Linear(in_features, out_features, bias=False)
        self.bn = nn.BatchNorm1d(out_features)
    def forward(self, x):
        x = self.linear(x)
        return self.bn(x)
net = LinearBN(32, 32).cuda().eval()  
print("weight: ", net.linear.weight.dtype) #torch.float32
img = torch.rand((1, 32)).cuda() #could also replace with 3d tensor, torch.rand((1, 32, 32))
with torch.cuda.amp.autocast(True):
    net = torch.jit.script(net)
    net = torch.jit.freeze(net)
lin_node = net.graph.findNode("aten::linear", True)
print("weight: ", lin_node.namedInput("weight").type().dtype()) #torch.float16
print("bias: ", lin_node.namedInput("bias").type().dtype()) #torch.float32
net(img) #weight, bias dtype inconsistency catched
```

### Versions

master `2bda2baad787923b064c747e619e62a6af969940` + #86706

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel