Repository: pytorch/functorch (language: python)

Issue:
autograd.Function is silently incorrect on functorch transforms
We should fix it. Not quite sure how though.

https://gist.github.com/zou3519/73a6189e21561f6ef5b42874e8a4826f


Discussion / hints:

Fix patch:
diff --git a/functorch/csrc/DynamicLayer.cpp b/functorch/csrc/DynamicLayer.cpp
index 46fb62660..ff1af1eec 100644
--- a/functorch/csrc/DynamicLayer.cpp
+++ b/functorch/csrc/DynamicLayer.cpp
@@ -104,7 +104,9 @@ class FuncTorchTLS : public FuncTorchTLSBase {
   }
 
   int64_t checkSupportsAutogradFunction() const override {
-    // Does nothing
+    TORCH_CHECK(dynamicLayerStack.size() <= 1, // we're inside a transform if the stack has more than the inital layer
+        "functorch functions (vmap, grad, vjp, etc.) currently do not support the use of autograd.Function. ",
+        "Please rewrite your function to not use autograd.Function while we work on fixing this");
     return 0;
   }
 


Test patch:
diff --git a/test/test_vmap.py b/test/test_vmap.py
index bb6e7d43e..885339092 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -34,7 +34,7 @@
 from collections import namedtuple
 
 import functorch
-from functorch import vmap, grad
+from functorch import vmap, grad, grad_and_value, jvp, vjp, jacrev, jacfwd
 from functorch._C import reshape_dim_into, reshape_dim_outof
 from functorch._src.make_functional import functional_init_with_buffers
 
@@ -3632,6 +3632,38 @@ def test_distributions(self, device, randomness, use_generator, batched_input):
                     assert torch.allclose(vmap_result[i], expected)
 
 
+class TestTransformFailure(TestCase):
+    @parametrize('transform', [vmap, grad, grad_and_value, vjp, jvp, jacrev, jacfwd])
+    def test_fails_with_autograd_function(self, device, transform):
+        class Test(torch.autograd.Function):
+            @staticmethod
+            def forward(_, input):
+                return input
+
+            @staticmethod
+            def backward(_, grad_input):
+                return grad_input
+
+        def f(x):
+            return Test.apply(x)
+
+        if transform == grad or transform == grad_and_value:
+            input = torch.tensor(4.)
+        else:
+            input = torch.randn(5)
+
+        if transform == vjp:
+            transform = functools.partial(transform, f)
+        elif transform == jvp:
+            input = (input,)
+            transform = functools.partial(transform, f, input)
+        else:
+            transform = transform(f)
+
+        with self.assertRaisesRegex(RuntimeError, "autograd.Function"):
+            transform(input)
+
+
 only_for = ("cpu", "cuda")
 instantiate_device_type_tests(TestVmapOperatorsOpInfo, globals(), only_for=only_for)
 
@@ -3640,6 +3672,7 @@ def test_distributions(self, device, randomness, use_generator, batched_input):
     globals(),
     only_for=only_for,
 )
+instantiate_device_type_tests(TestTransformFailure, globals(), only_for=only_for)
 
 if __name__ == '__main__':
     run_tests()