don't specialize when grad tracking tensors are activated (#140828)

bobrenjc93 · pytorchmergebot · commit 90051560048f · 2024-11-17T18:28:47.000Z
Fixes `python test/dynamo/test_inline_inbuilt_nn_modules.py InlineInbuiltNNModulesFuncTorchHigherOrderOpTests.test_grad_non_tensor_input_inline_inbuilt_nn_modules` when `specialize_float=False` Pull Request resolved: #140828 Approved by: https://github.com/ezyang ghstack dependencies: #140830, #140832
diff --git a/torch/_dynamo/variables/builder.py b/torch/_dynamo/variables/builder.py
@@ -1914,6 +1914,14 @@ def wrap_symfloat(self, value):
         # time.
 
         wrapped_value = torch.tensor(value, dtype=torch.float64)
+
+        # We don't support specializing floats for grad checking tensors
+        # See https://github.com/pytorch/pytorch/pull/140828 for more
+        # context.
+        if torch._C._functorch.is_gradtrackingtensor(wrapped_value):
+            self.install_guards(GuardBuilder.CONSTANT_MATCH)
+            return ConstantVariable.create(value=value, source=self.source)
+
         # TODO: Switch RandomValueSource over to use this, this is more
         # accurate
         assert not isinstance(self.get_source(), RandomValueSource)
