[JIT] Compilation-induced discrepancy in F.instance_norm when passing input as running stats

username: zhouxiaoyaozzz
cooked: <h3><a name="p-470408-bug-description-1" class="anchor" href="#p-470408-bug-description-1"></a>Bug Description</h3>
<p>When scripting a model containing <code>F.instance_norm</code> with broadcasted input, JIT-compiled results differ from eager mode.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F


class NeuralModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.mlist = nn.ParameterList([
            nn.Parameter(torch.tensor([[
                [[5.0269, 3.4145, 3.8807],
                 [6.3197, 6.5815, 5.3826],
                 [3.4846, 4.3035, 6.9038]],

                [[3.7721, 3.0422, 4.3315],
                 [3.2518, 3.5617, 6.4604],
                 [3.5747, 5.1481, 6.6348]],

                [[3.0173, 3.5291, 6.8552],
                 [5.6125, 6.2321, 6.3142],
                 [4.7299, 4.2638, 5.2731]]
            ]], dtype=torch.float64))
        ])

    def forward(self, x):
        expanded = x.unsqueeze(-1).unsqueeze(-1)  # (3,) -&gt; (3,1,1)
        multiplied = x * self.mlist[0]  # Broadcast multiply
        inst_norm = F.instance_norm(multiplied, x, x)
        log_softmax = F.log_softmax(multiplied, dim=-1)
        bilinear = F.interpolate(log_softmax, scale_factor=1.0, mode='bilinear')  # 修复：1.0 而不是 1

        return {
            'v0_0': expanded,
            'v6_0': inst_norm,
            'v2_0': bilinear
        }

input_data = torch.tensor([494.91649119, 528.01665228, 492.01463052], dtype=torch.float64)
model = NeuralModel()
with torch.no_grad():
    output_eager = model(input_data)

model_scripted = torch.jit.script(model)
with torch.no_grad():
    output_scripted = model_scripted(input_data)

print("NONJIT (v6_0):", output_eager['v6_0'])
print("JIT  (v6_0):", output_scripted['v6_0'])
print("consistency:", torch.allclose(output_eager['v6_0'], output_scripted['v6_0']))
</code></pre>
<h3><a name="p-470408-output-2" class="anchor" href="#p-470408-output-2"></a>output:</h3>
<pre><code class="lang-auto">NONJIT (v6_0): tensor([[[[-0.0792, -1.1555, -0.9882],
          [ 0.9260,  1.4719,  0.1728],
          [-1.2785, -0.4180,  1.3487]],

         [[-0.5777, -0.9980, -0.1514],
          [-0.9931, -0.5555,  1.5384],
          [-0.7353,  0.7958,  1.6768]],

         [[-1.7688, -1.1584,  1.3313],
          [ 0.3497,  1.1958,  0.8923],
          [-0.3708, -0.5185,  0.0474]]]], dtype=torch.float64)
JIT  (v6_0): tensor([[[[ 0.1992, -1.2898, -1.0547],
          [ 1.2781,  1.2202,  0.0836],
          [-1.0879, -0.5852,  1.2365]],

         [[-0.3739, -1.1909, -0.2243],
          [-0.8553, -0.7344,  1.5646],
          [-0.5565,  0.6596,  1.7112]],

         [[-1.6616, -1.3545,  1.2870],
          [ 0.7235,  1.0046,  0.8354],
          [-0.0877, -0.7133, -0.0335]]]], dtype=torch.float64)
consistency: False

</code></pre>
<h3><a name="p-470408-versions-3" class="anchor" href="#p-470408-versions-3"></a>Versions</h3>
<p>Collecting environment information…<br>
PyTorch version: 2.0.1+cpu<br>
Is debug build: False<br>
CUDA used to build PyTorch: Could not collect<br>
ROCM used to build PyTorch: N/A</p>
<p>OS: Could not collect<br>
GCC version: Could not collect<br>
Clang version: 20.1.2<br>
CMake version: version 4.0.0<br>
Libc version: N/A</p>
<p>Python version: 3.9.7 (tags/v3.9.7:1016ef3, Aug 30 2021, 20:19:38) [MSC v.1929 64 bit (AMD64)] (64-bit runtime)<br>
Python platform: Windows-10-10.0.26100-SP0<br>
Is CUDA available: False<br>
CUDA runtime version: Could not collect<br>
CUDA_MODULE_LOADING set to: N/A<br>
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 4060 Laptop GPU<br>
Nvidia driver version: 560.94<br>
cuDNN version: Could not collect<br>
HIP runtime version: N/A<br>
MIOpen runtime version: N/A<br>
Is XNNPACK available: True</p>
<p>Versions of relevant libraries:<br>
[pip3] numpy==1.26.1<br>
[pip3] torch==2.0.1<br>
[conda] Could not collect</p>