Revert "Support setting grad_dtype on leaf tensors (#162815)"

This reverts commit dca7398.

Reverted #162815 on behalf of https://github.com/yangw-dev due to break internal test D83850533, see more details below ([comment](#162815 (comment)))

Author: pytorchmergebot

aten/src/ATen/core/Tensor.cpp

@@ -173,12 +173,4 @@ unsigned TensorBase::_register_hook(std::function<TensorBase(const TensorBase&)>
   return impl::GetVariableHooks()->_register_hook(*this, std::move(hook));
 }
 
-std::optional<ScalarType> TensorBase::grad_dtype() const {
-  return impl::GetVariableHooks()->grad_dtype(*this);
-}
-
-void TensorBase::set_grad_dtype(const std::optional<ScalarType>& grad_dtype) const {
-  return impl::GetVariableHooks()->set_grad_dtype(*this, grad_dtype);
-}
-
 } // namespace at

aten/src/ATen/core/TensorBase.h

@@ -930,10 +930,6 @@ class TORCH_API TensorBase {
 
   const TensorBase& requires_grad_(bool _requires_grad=true) const;
 
-  std::optional<ScalarType> grad_dtype() const;
-
-  void set_grad_dtype(const std::optional<ScalarType>& grad_dtype) const;
-
   // View Variables
   //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

aten/src/ATen/core/VariableHooksInterface.h

@@ -68,8 +68,6 @@ struct TORCH_API VariableHooksInterface {
       const c10::OperatorHandle& op,
       c10::DispatchKeySet dispatch_keys,
       torch::jit::Stack* stack) const = 0;
-  virtual std::optional<c10::ScalarType> grad_dtype(const TensorBase&) const = 0;
-  virtual void set_grad_dtype(const TensorBase&, const std::optional<c10::ScalarType>&) const = 0;
 };
 
 TORCH_API void SetVariableHooks(VariableHooksInterface* hooks);

test/dynamo/test_backward_higher_order_ops.py

@@ -140,7 +140,7 @@ def forward(self, L_inputs_ : list, s69: "Sym(s21)", L_sizes_0_: "f32[0, s21]"):
 
         size: "Sym(s21)" = l_sizes_0_.size(1);  l_sizes_0_ = None
 
-        validate_outputs = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem], [((None, None, device(type='cpu'), 6, 0, None), [size], False, 6)]);  getitem = size = None
+        validate_outputs = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem], [((None, None, device(type='cpu'), 6, 0, None), [size], False)]);  getitem = size = None
         getitem_9: "f32[s21]" = validate_outputs[0];  validate_outputs = None
 
         call_aot_bwd_prologue = torch__dynamo_compiled_autograd_call_aot_bwd_prologue((), [], getitem_9);  getitem_9 = None
@@ -171,7 +171,7 @@ def forward(self, L_inputs_ : list, s69: "Sym(s21)", L_sizes_0_: "f32[0, s21]"):
 
         size: "Sym(s21)" = l_sizes_0_.size(1);  l_sizes_0_ = None
 
-        validate_outputs = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem], [((None, None, device(type='cpu'), 6, 0, None), [size], False, 6)]);  getitem = size = None
+        validate_outputs = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem], [((None, None, device(type='cpu'), 6, 0, None), [size], False)]);  getitem = size = None
         getitem_9: "f32[s21]" = validate_outputs[0];  validate_outputs = None
 
         call_aot_bwd_prologue = torch__dynamo_compiled_autograd_call_aot_bwd_prologue((), [], getitem_9);  getitem_9 = None
@@ -255,7 +255,7 @@ def forward(self, L_inputs_ : list, s69: "Sym(s21)", L_sizes_0_: "f32[0, s21]",
 
         size: "Sym(s21)" = l_sizes_0_.size(1);  l_sizes_0_ = None
 
-        validate_outputs = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem], [((None, None, device(type='cpu'), 6, 0, None), [size], False, 6)]);  getitem = size = None
+        validate_outputs = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem], [((None, None, device(type='cpu'), 6, 0, None), [size], False)]);  getitem = size = None
         getitem_9: "f32[s21]" = validate_outputs[0];  validate_outputs = None
 
         call_aot_bwd_prologue = torch__dynamo_compiled_autograd_call_aot_bwd_prologue((), [], getitem_9);  getitem_9 = None

test/inductor/test_compiled_autograd.py

@@ -3604,12 +3604,12 @@ def forward(self, inputs, sizes, scalars, hooks, packed_data):
         unwrap_maybe_dynamic_int_18 = torch__dynamo_external_utils_unwrap_maybe_dynamic_int(getitem_23);  getitem_23 = None
         unwrap_maybe_dynamic_int_19 = torch__dynamo_external_utils_unwrap_maybe_dynamic_int(getitem_24);  getitem_24 = None
 
-        validate_outputs = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem], [((None, None, device(type='cpu'), 6, 0, None), [], True, 6)]);  getitem = None
+        validate_outputs = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem], [((None, None, device(type='cpu'), 6, 0, None), [], True)]);  getitem = None
         getitem_25 = validate_outputs[0];  validate_outputs = None
 
         sum_backward0 = torch__dynamo_compiled_autograd_ops_SumBackward0([getitem_25], [True], [unwrap_maybe_dynamic_int, unwrap_maybe_dynamic_int_1]);  getitem_25 = unwrap_maybe_dynamic_int = unwrap_maybe_dynamic_int_1 = None
         getitem_26 = sum_backward0[0];  sum_backward0 = None
-        validate_outputs_1 = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem_26], [((None, None, device(type='cpu'), 6, 0, None), [unwrap_maybe_dynamic_int_2, unwrap_maybe_dynamic_int_3], True, 6)]);  getitem_26 = unwrap_maybe_dynamic_int_2 = unwrap_maybe_dynamic_int_3 = None
+        validate_outputs_1 = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem_26], [((None, None, device(type='cpu'), 6, 0, None), [unwrap_maybe_dynamic_int_2, unwrap_maybe_dynamic_int_3], True)]);  getitem_26 = unwrap_maybe_dynamic_int_2 = unwrap_maybe_dynamic_int_3 = None
         getitem_27 = validate_outputs_1[0];  validate_outputs_1 = None
 
         getitem_28 = hooks[0];  getitem_28 = None
@@ -3631,7 +3631,7 @@ def forward(self, inputs, sizes, scalars, hooks, packed_data):
         call_backward = torch__dynamo_external_utils_call_backward(getitem_33, (), make_subclass);  getitem_33 = make_subclass = None
         getitem_36 = call_backward[0]
         getitem_37 = call_backward[1];  call_backward = None
-        validate_outputs_2 = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem_36, getitem_37], [((None, None, device(type='cpu'), 6, 0, None), [unwrap_maybe_dynamic_int_16, unwrap_maybe_dynamic_int_17], False, 6), ((None, None, device(type='cpu'), 6, 0, None), [unwrap_maybe_dynamic_int_18, unwrap_maybe_dynamic_int_19], False, 6)]);  getitem_36 = getitem_37 = unwrap_maybe_dynamic_int_16 = unwrap_maybe_dynamic_int_17 = unwrap_maybe_dynamic_int_18 = unwrap_maybe_dynamic_int_19 = None
+        validate_outputs_2 = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem_36, getitem_37], [((None, None, device(type='cpu'), 6, 0, None), [unwrap_maybe_dynamic_int_16, unwrap_maybe_dynamic_int_17], False), ((None, None, device(type='cpu'), 6, 0, None), [unwrap_maybe_dynamic_int_18, unwrap_maybe_dynamic_int_19], False)]);  getitem_36 = getitem_37 = unwrap_maybe_dynamic_int_16 = unwrap_maybe_dynamic_int_17 = unwrap_maybe_dynamic_int_18 = unwrap_maybe_dynamic_int_19 = None
         getitem_39 = validate_outputs_2[0]
 
         call_accumulate_grad_1 = torch__dynamo_external_utils_call_accumulate_grad(getitem_4, getitem_39, False);  getitem_4 = getitem_39 = call_accumulate_grad_1 = None
@@ -3866,12 +3866,12 @@ def forward(self, inputs, sizes, scalars, hooks, packed_data):
         unwrap_maybe_dynamic_int_10 = torch__dynamo_external_utils_unwrap_maybe_dynamic_int(getitem_12);  getitem_12 = None
         unwrap_maybe_dynamic_int_11 = torch__dynamo_external_utils_unwrap_maybe_dynamic_int(getitem_13);  getitem_13 = None
 
-        validate_outputs = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem], [((None, None, device(type='cpu'), 6, 0, None), [], False, 6)]);  getitem = None
+        validate_outputs = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem], [((None, None, device(type='cpu'), 6, 0, None), [], False)]);  getitem = None
         getitem_14 = validate_outputs[0];  validate_outputs = None
 
         sum_backward0 = torch__dynamo_compiled_autograd_ops_SumBackward0([getitem_14], [True], [unwrap_maybe_dynamic_int, unwrap_maybe_dynamic_int_1]);  getitem_14 = unwrap_maybe_dynamic_int = unwrap_maybe_dynamic_int_1 = None
         getitem_15 = sum_backward0[0];  sum_backward0 = None
-        validate_outputs_1 = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem_15], [((None, None, device(type='cpu'), 6, 0, None), [unwrap_maybe_dynamic_int_2, unwrap_maybe_dynamic_int_3], False, 6)]);  getitem_15 = unwrap_maybe_dynamic_int_2 = unwrap_maybe_dynamic_int_3 = None
+        validate_outputs_1 = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem_15], [((None, None, device(type='cpu'), 6, 0, None), [unwrap_maybe_dynamic_int_2, unwrap_maybe_dynamic_int_3], False)]);  getitem_15 = unwrap_maybe_dynamic_int_2 = unwrap_maybe_dynamic_int_3 = None
         getitem_16 = validate_outputs_1[0];  validate_outputs_1 = None
 
         getitem_17 = hooks[0]
@@ -3883,7 +3883,7 @@ def forward(self, inputs, sizes, scalars, hooks, packed_data):
         mul_backward0 = torch__dynamo_compiled_autograd_ops_MulBackward0([getitem_16], [True, True], call_hook, 6, call_hook_1, 6);  getitem_16 = call_hook = call_hook_1 = None
         getitem_21 = mul_backward0[0]
         getitem_22 = mul_backward0[1];  mul_backward0 = None
-        validate_outputs_2 = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem_21, getitem_22], [((None, None, device(type='cpu'), 6, 0, None), [unwrap_maybe_dynamic_int_4, unwrap_maybe_dynamic_int_5], False, 6), ((None, None, device(type='cpu'), 6, 0, None), [unwrap_maybe_dynamic_int_6, unwrap_maybe_dynamic_int_7], False, 6)]);  getitem_21 = getitem_22 = unwrap_maybe_dynamic_int_4 = unwrap_maybe_dynamic_int_5 = unwrap_maybe_dynamic_int_6 = unwrap_maybe_dynamic_int_7 = None
+        validate_outputs_2 = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem_21, getitem_22], [((None, None, device(type='cpu'), 6, 0, None), [unwrap_maybe_dynamic_int_4, unwrap_maybe_dynamic_int_5], False), ((None, None, device(type='cpu'), 6, 0, None), [unwrap_maybe_dynamic_int_6, unwrap_maybe_dynamic_int_7], False)]);  getitem_21 = getitem_22 = unwrap_maybe_dynamic_int_4 = unwrap_maybe_dynamic_int_5 = unwrap_maybe_dynamic_int_6 = unwrap_maybe_dynamic_int_7 = None
         getitem_23 = validate_outputs_2[0]
         getitem_24 = validate_outputs_2[1];  validate_outputs_2 = None
 
@@ -3892,7 +3892,7 @@ def forward(self, inputs, sizes, scalars, hooks, packed_data):
         call_hook_2 = torch__dynamo_external_utils_call_hook(getitem_25, getitem_26, hook_type = 'unpack_hook');  getitem_25 = getitem_26 = None
         cos_backward0 = torch__dynamo_compiled_autograd_ops_CosBackward0([getitem_24], [True], call_hook_2);  getitem_24 = call_hook_2 = None
         getitem_27 = cos_backward0[0];  cos_backward0 = None
-        validate_outputs_3 = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem_27], [((None, None, device(type='cpu'), 6, 0, None), [unwrap_maybe_dynamic_int_8, unwrap_maybe_dynamic_int_9], False, 6)]);  getitem_27 = unwrap_maybe_dynamic_int_8 = unwrap_maybe_dynamic_int_9 = None
+        validate_outputs_3 = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem_27], [((None, None, device(type='cpu'), 6, 0, None), [unwrap_maybe_dynamic_int_8, unwrap_maybe_dynamic_int_9], False)]);  getitem_27 = unwrap_maybe_dynamic_int_8 = unwrap_maybe_dynamic_int_9 = None
         getitem_28 = validate_outputs_3[0];  validate_outputs_3 = None
         add = torch.add(getitem_23, getitem_28);  getitem_23 = getitem_28 = None
 
@@ -3901,7 +3901,7 @@ def forward(self, inputs, sizes, scalars, hooks, packed_data):
         call_hook_3 = torch__dynamo_external_utils_call_hook(getitem_29, getitem_30, hook_type = 'unpack_hook');  getitem_29 = getitem_30 = None
         sin_backward0 = torch__dynamo_compiled_autograd_ops_SinBackward0([add], [True], call_hook_3);  add = call_hook_3 = None
         getitem_31 = sin_backward0[0];  sin_backward0 = None
-        validate_outputs_4 = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem_31], [((None, None, device(type='cpu'), 6, 0, None), [unwrap_maybe_dynamic_int_10, unwrap_maybe_dynamic_int_11], False, 6)]);  getitem_31 = unwrap_maybe_dynamic_int_10 = unwrap_maybe_dynamic_int_11 = None
+        validate_outputs_4 = torch__dynamo_compiled_autograd_ops_validate_outputs([getitem_31], [((None, None, device(type='cpu'), 6, 0, None), [unwrap_maybe_dynamic_int_10, unwrap_maybe_dynamic_int_11], False)]);  getitem_31 = unwrap_maybe_dynamic_int_10 = unwrap_maybe_dynamic_int_11 = None
         getitem_32 = validate_outputs_4[0];  validate_outputs_4 = None
 
         call_accumulate_grad = torch__dynamo_external_utils_call_accumulate_grad(getitem_1, getitem_32, False);  getitem_1 = getitem_32 = call_accumulate_grad = None
@@ -5266,7 +5266,6 @@ def wrap_test_class(orig_cls):
         "test_dropout_inductor",  # functionalize_rng_ops not yet supported
         "test_function_with_kwargs",  # functionalize_rng_ops not yet supported
         "test_module",  # functionalize_rng_ops not yet supported
-        "test_grad_dtype",  # AttributeError: args / Float did not match Double
     },
     "eager": {  # will be run without torch.compiling the CA graph
         "test_setup_context_when_forward_has_default_args",  # autograd.Function with class methods

test/jit/test_builtins.py

@@ -326,8 +326,6 @@ def fn(x):
             # This has a longer implementation, maybe not worth copying to
             # TorchScript if named tensors don't work there anyways
             "names",
-            # We don't plan to support grad_dtype in TorchScript
-            "grad_dtype",
         }
 
         for p in properties:

test/test_autograd.py

@@ -3694,130 +3694,6 @@ def test_sparse_gather_x_scalar(self):
     def test_sparse_gather_both_scalar(self):
         self._test_sparse_gather((), (), 0)
 
-    @skipIfTorchDynamo("grad_dtype not supported in compile")
-    def test_grad_dtype(self):
-        leaf = torch.tensor([1.0, 2.0], requires_grad=True)
-        # Default to tensor's dtype
-        self.assertEqual(leaf.grad_dtype, torch.float32)
-        leaf.grad_dtype = torch.float16
-        self.assertEqual(leaf.grad_dtype, torch.float16)
-        leaf.grad_dtype = None  # Allow any dtype
-        self.assertIsNone(leaf.grad_dtype)
-
-        # get/set grad_dtype is only allowed on leaf tensors
-        non_leaf = leaf * 2
-        self.assertFalse(non_leaf.is_leaf)
-        with self.assertRaisesRegex(
-            RuntimeError, "grad_dtype can only be accessed on leaf tensors"
-        ):
-            _ = non_leaf.grad_dtype
-        with self.assertRaisesRegex(
-            RuntimeError, "grad_dtype can only be set on leaf tensors"
-        ):
-            non_leaf.grad_dtype = torch.float16
-
-        # Manual setting
-        x = torch.tensor([1.0, 2.0], requires_grad=True)
-        grad_match = torch.tensor([1.0, 1.0])
-        x.grad = grad_match
-        self.assertEqual(x.grad.dtype, torch.float32)
-
-        x.grad = None
-        x.grad_dtype = torch.float16
-        grad_mismatch = torch.tensor([1.0, 1.0])
-        with self.assertRaisesRegex(
-            RuntimeError,
-            "attempting to assign a gradient with dtype.*float.*to a tensor with grad_dtype.*Half",
-        ):
-            x.grad = grad_mismatch
-
-        # When grad_dtype is None, any dtype is allowed
-        x.grad = None
-        x.grad_dtype = None
-        grad_any = torch.tensor([1.0, 1.0], dtype=torch.float64)
-        x.grad = grad_any
-        self.assertEqual(x.grad.dtype, torch.float64)
-
-        # Incoming gradient case
-        class MismatchedGradientFunction(torch.autograd.Function):
-            @staticmethod
-            def forward(ctx, inp):
-                return inp * 2
-
-            @staticmethod
-            def backward(ctx, grad_output):
-                return grad_output.to(torch.float64)
-
-        d = torch.tensor([1.0, 2.0], requires_grad=True)
-        output = MismatchedGradientFunction.apply(d)
-        loss = output.sum()
-        loss.backward()
-        # Default behavior is to cast to tensor dtype
-        self.assertEqual(d.grad.dtype, torch.float32)
-        self.assertTrue(torch.allclose(d.grad, torch.tensor([1.0, 1.0])))
-
-        e = torch.tensor([3.0, 4.0], requires_grad=True)
-        e.grad_dtype = None
-        output_e = MismatchedGradientFunction.apply(e)
-        loss_e = output_e.sum()
-        loss_e.backward()
-        # No casting is done if set to None.
-        self.assertTrue(
-            torch.allclose(e.grad, torch.tensor([1.0, 1.0], dtype=torch.float64))
-        )
-
-        f = torch.tensor([5.0, 6.0], requires_grad=True)
-        f.grad_dtype = torch.float16  # Expect float16 gradients
-        output_f = MismatchedGradientFunction.apply(f)
-        loss_f = output_f.sum()
-        loss_f.backward()
-        self.assertTrue(
-            torch.allclose(f.grad, torch.tensor([1.0, 1.0], dtype=torch.float16))
-        )
-
-        # Setting grad_dtype when gradient already exists
-        g = torch.tensor([1.0, 2.0], requires_grad=True)
-        g.grad = torch.tensor([1.0, 1.0])
-        g.grad_dtype = torch.float32
-        self.assertEqual(g.grad_dtype, torch.float32)
-        with self.assertRaisesRegex(
-            RuntimeError, "Cannot set grad_dtype.*because there is already a gradient"
-        ):
-            g.grad_dtype = torch.float16
-        g.grad_dtype = None
-        self.assertIsNone(g.grad_dtype)
-        g.grad = None
-        g.grad_dtype = torch.float16
-        self.assertEqual(g.grad_dtype, torch.float16)
-
-        # Test the case where there is an existing accumulate grad
-        h = torch.tensor([1.0, 2.0], requires_grad=True)
-        _ = h.clone()
-        h.grad_dtype = None
-        output = MismatchedGradientFunction.apply(h)
-        output.sum().backward()
-        self.assertEqual(h.grad.dtype, torch.float64)
-
-        # Mixed accumulation cases
-        k = torch.tensor([1.0, 2.0], requires_grad=True)
-        k.grad_dtype = None
-        y = k * 2
-        y.sum().backward()
-        k.grad = k.grad.to(torch.bfloat16)
-        y2 = k * 3
-        # Doesn't type promote to float32, always coerce to current .grad's dtype.
-        # This is because the accumulation is done in-place on the existing grad.
-        self.assertEqual(k.grad.dtype, torch.bfloat16)
-
-        l = torch.tensor([3.0, 4.0], requires_grad=True, dtype=torch.bfloat16)
-        l.grad_dtype = None
-        z = l * 2
-        z.sum().backward()
-        l.grad = l.grad.to(torch.float32)
-        z2 = l * 3
-        z2.sum().backward()
-        self.assertEqual(l.grad.dtype, torch.float32)
-
     def test_gc_in_destructor(self):
         """
         Previously, if a Function destructor triggered a garbage collection,

torch/_C/__init__.pyi.in

@@ -1917,7 +1917,6 @@ class TensorBase(metaclass=_TensorMeta):
     names: list[str]
     device: _device
     dtype: _dtype
-    grad_dtype: _dtype | None
     layout: _layout
     real: Tensor
     imag: Tensor

torch/_tensor_docs.py

@@ -6624,36 +6624,6 @@ def callable(a, b) -> number
 """,
 )
 
-add_docstr_all(
-    "grad_dtype",
-    r"""
-The allowed dtype of :attr:``grad`` for this tensor.
-
-:attr:``grad_dtype`` can be set to a specific dtype or ``None``. By default,
-``t.grad_dtype == t.dtype``. When not None, the autograd engine casts
-incoming gradients to this dtype. This attribute is only accessible and
-settable for leaf tensors.
-
-.. warning::
-    Use with caution. Diverging the dtypes of a tensor and its gradient may
-    break downstream systems that assume they match.
-
-Example::
-
-    >>> x = torch.tensor([1.0, 2.0], requires_grad=True)
-    >>> x.grad_dtype
-    torch.float32
-
-    >>> x.grad_dtype = torch.float16
-    >>> x.grad_dtype
-    torch.float16
-
-    >>> # Allow any gradient dtype
-    >>> x.grad_dtype = None
-    >>> x.grad_dtype
-""",
-)
-
 add_docstr_all(
     "retain_grad",
     r"""