this is a super minor nit - but from what I can tell the contents of fold_linear_bn.(h|cpp) are essentially implementation details of frozen_linear_folding.cpp, is that correct? if so we can probably move this into the anonymous namespace of frozen_linear_folding.cpp right?

(this is mostly just a matter of preference tbh)

That's correct. Actually, fold_linear_bn.h is mostly copied from fold_conv_bn.h, so I thought I would follow it's style. I'm fine with either options -- what do you think?

I guess it's fine as it is

this looks like it's mostly copied from frozen_conv_folding.cpp, is that accurate? recently we had some issues with autocasting, see #77617. Can you add some tests like this to confirm that it's not an issue in this case?

this looks like it's mostly copied from frozen_conv_folding.cpp, is that accurate?

That's correct, this is mostly copied from frozen_conv_folding.cpp.

recently we had some issues with autocasting, see #77617

Thanks for the notice. I verified that similar issue occurs for jit autocasting for lin-bn1d only and on cuda only (inputs are only casted to half on cuda).

The dtype mistmatch is checked in addmm with {2,3}d input tensor, which is the correct usage of bn1d; however, the issue will also occur for scripting (not tracing) with an incorrect usage of bn{2,3}d with {2,3}d input tensor (incorrect usage, expected input dims are 4d, 5d for bn2d, and bn3d respectively). Below, replace nn.BatchNorm1d(out_features) with nn.BatchNorm{2,3}d(out_features).

To reproduce RuntimeError: self and mat2 must have the same dtype:

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.BatchNorm1d(out_features)
    
    def forward(self, x):
        x = self.linear(x)
        return self.bn(x)

net = LinearBN(32, 32).cuda().eval()  
print(net.linear.weight.dtype) #torch.float32
img = torch.rand((1, 32, 32)).cuda()
with torch.cuda.amp.autocast(True):
    net = torch.jit.script(net)
    net = torch.jit.freeze(net)
print(net.graph)
lin_node = net.graph.findNode("aten::linear", True)
print(lin_node.namedInput("weight").type().dtype()) #torch.float16
print(lin_node.namedInput("bias").type().dtype()) #torch.float32
net(img)