Repository: pytorch/functorch (language: python)

Issue:
Slight numerical differences when using `linear` with functorch transforms
It looks like functorch has its own copy of the C++ decomposition for `linear`, for performance reasons: https://github.com/pytorch/functorch/blob/main/functorch/csrc/PyTorchOperatorHacks.cpp

It looks like the implementations have diverged over time, leading to numeric differences: https://github.com/pytorch/pytorch/blob/56bed0dcfe7ca9047e5c95a6f3d7fcb0ec403b0c/aten/src/ATen/native/Linear.cpp#L37.

You can see an example below, where when I run `linear` with and without `functionalize()`, I get an `addmm` in one case and an `mm` + `add` in another case, leading to a slight numeric difference.

Should there be a better long term plan for keeping these in sync? Maybe a numeric accuracy test for the ops that are copied over?

===== code =====

import torch
from functorch import make_fx
from functorch.experimental import functionalize


def f(x, y, z):
    return torch._C._nn.linear(x, y, z)


x = torch.randn(14, 1, 384)
y = torch.randn(96, 384)
z = torch.randn(96)

out = make_fx(f)(x, y, z)
print(out.graph)

out = make_fx(functionalize(f))(x, y, z)
print(out.graph)


a = f(x, y, z)
b = functionalize(f)(x, y, z)
print(f"max difference: {torch.max(torch.abs(a - b))}")
print(torch.allclose(a, b))


==== outputs ====

graph():
    %x_1 : [#users=1] = placeholder[target=x_1]
    %y_1 : [#users=1] = placeholder[target=y_1]
    %z_1 : [#users=1] = placeholder[target=z_1]
    %t : [#users=1] = call_function[target=torch.ops.aten.t](args = (%y_1,), kwargs = {})
    %view : [#users=1] = call_function[target=torch.ops.aten.view](args = (%x_1, [14, 384]), kwargs = {})
    %addmm : [#users=1] = call_function[target=torch.ops.aten.addmm](args = (%z_1, %view, %t), kwargs = {})
    %view_1 : [#users=1] = call_function[target=torch.ops.aten.view](args = (%addmm, [14, 1, 96]), kwargs = {})
    return view_1
graph():
    %x_1 : [#users=1] = placeholder[target=x_1]
    %y_1 : [#users=1] = placeholder[target=y_1]
    %z_1 : [#users=1] = placeholder[target=z_1]
    %t : [#users=1] = call_function[target=torch.ops.aten.t](args = (%y_1,), kwargs = {})
    %view : [#users=1] = call_function[target=torch.ops.aten.view](args = (%x_1, [14, 384]), kwargs = {})
    %mm : [#users=1] = call_function[target=torch.ops.aten.mm](args = (%view, %t), kwargs = {})
    %_unsafe_view : [#users=1] = call_function[target=torch.ops.aten._unsafe_view](args = (%mm, [14, 1, 96]), kwargs = {})
    %add : [#users=1] = call_function[target=torch.ops.aten.add](args = (%_unsafe_view, %z_1), kwargs = {})
    return add
max difference: 3.814697265625e-06
False


Discussion / hints:

Fix patch:
diff --git a/functorch/csrc/PyTorchOperatorHacks.cpp b/functorch/csrc/PyTorchOperatorHacks.cpp
index b8e7079e6..d72702d2e 100644
--- a/functorch/csrc/PyTorchOperatorHacks.cpp
+++ b/functorch/csrc/PyTorchOperatorHacks.cpp
@@ -118,6 +118,12 @@ Tensor linear_hack(const Tensor& input, const Tensor& weight, const c10::optiona
     // Fused op is marginally faster.
     return at::addmm(*bias, input, weight.t());
   }
+  if (input.dim() == 3 && bias->defined() && input.is_contiguous()) {
+    // Also hit the fused path for contiguous 3D input.
+    const auto input_sizes = input.sizes();
+    const auto result = at::addmm(*bias, input.view({input_sizes[0] * input_sizes[1], input_sizes[2]}), weight.t());
+    return result.view({input_sizes[0], input_sizes[1], result.size(1)});
+  }
   auto output = at::matmul(input, weight.t());
   if (bias->defined()) {
     const auto& stack = getDynamicLayerStack();


Test patch:
diff --git a/test/test_eager_transforms.py b/test/test_eager_transforms.py
index a59956df1..22e1b0ec4 100644
--- a/test/test_eager_transforms.py
+++ b/test/test_eager_transforms.py
@@ -2858,6 +2858,20 @@ def f(x: torch.Tensor) -> torch.Tensor:
             return y
         self._check_functionalize_correctness(f, torch.zeros(4, 2, device=device))
 
+    # See https://github.com/pytorch/functorch/issues/780
+    def test_linear(self, device):
+
+        def f(x, y, z) -> torch.Tensor:
+            return torch._C._nn.linear(x, y, z)
+
+        x = torch.randn(14, 1, 384, device=device)
+        y = torch.randn(96, 384, device=device)
+        z = torch.randn(96, device=device)
+
+        out_expected = f(x, y, z)
+        out_actual = functionalize(f)(x, y, z)
+        self.assertEqual(out_expected, out_actual)
+
     def test_multioutput_inplace_slice_view(self, device):
 
         def f(x: torch.Tensor) -> torch.Tensor: