pytorch
diff --git a/‎test/test_autograd.py
+41 b/‎test/test_autograd.py
+41
diff --git a/‎torch/overrides.py
+12 b/‎torch/overrides.py
+12
diff --git a/‎torch/utils/checkpoint.py
+10-1 b/‎torch/utils/checkpoint.py
+10-1
@@ -80,6 +80,7 @@
     TestCase,
     xfailIfTorchDynamo,
 )
+from torch.overrides import TorchFunctionMode
 from torch.utils._mode_utils import no_dispatch
 from torch.utils._python_dispatch import TorchDispatchMode
 from torch.utils.checkpoint import (
@@ -7514,6 +7515,46 @@ def forward(self, dict_input):
         ):
             self.assertEqual(param.grad, checkpoint_param.grad)
 
+     @xfailIfTorchDynamo
+    def test_checkpointing_preserves_torch_function_mode_stack(self):
+        log = []
+
+        def get_mode_class(n):
+            class Func(TorchFunctionMode):
+                def __torch_function__(self, func, types, args, kwargs=None):
+                    kwargs = {} if kwargs is None else kwargs
+                    log.append(f"mode{n}")
+                    return func(*args, **kwargs)
+
+            return Func
+
+        Mode1 = get_mode_class(1)
+        Mode2 = get_mode_class(2)
+        Mode3 = get_mode_class(3)
+
+        def func(x):
+            return x.sin().cos()
+
+        def context_fn():
+            return Mode3(), Mode3()
+
+        with Mode1():
+            with Mode2():
+                a = torch.tensor(1.0, requires_grad=True)
+
+                log = []
+                out = checkpoint(func, a, use_reentrant=False, context_fn=context_fn)
+                self.assertTrue(
+                    log[-3] == "mode3" and log[-2] == "mode2" and log[-1] == "mode1"
+                )
+
+
+                log = []
+                out.backward()
+                self.assertTrue(
+                    log[-3] == "mode3" and log[-2] == "mode2" and log[1] == "mode1"
+                )
+
     def test_callback_adds_callback(self):
         called = [0]
 
 
@@ -2073,6 +2073,18 @@ def _pop_mode():
     return old
 
 
+@contextlib.contextmanager
+def _apply_torch_function_mode_stack(mode_stack):
+    stack_len = len(mode_stack)
+    try:
+        for mode in mode_stack:
+            _push_mode(mode)
+        yield
+    finally:
+        for _ in range(stack_len):
+            _pop_mode()
+
+
 @contextlib.contextmanager
 def _pop_mode_temporarily():
     old = _pop_mode()
 
@@ -15,6 +15,8 @@
 from torch.utils._pytree import tree_map
 from torch.testing._internal.logging_tensor import capture_logs, LoggingTensorMode
 from torch.utils._python_dispatch import TorchDispatchMode
+from torch.overrides import _get_current_function_mode_stack, _apply_torch_function_mode_stack
+
 
 __all__ = [
     "checkpoint",
@@ -1496,6 +1498,8 @@ def _checkpoint_without_reentrant_generator(
             had_device_in_fwd = True
             fwd_devices, fwd_device_states = get_device_states(*args)
 
+    torch_function_mode_stack = _get_current_function_mode_stack()
+
     def recompute_fn(*inputs):
         kwargs, *args = inputs
         # This will be called later during recomputation. This wrapping enables
@@ -1514,7 +1518,12 @@ def recompute_fn(*inputs):
             device_autocast_ctx = torch.amp.autocast(
                 device_type=device_type, **device_autocast_kwargs
             ) if torch.amp.is_autocast_available(device_type) else contextlib.nullcontext()
-            with device_autocast_ctx, torch.amp.autocast("cpu", **cpu_autocast_kwargs), recompute_context:  # type: ignore[attr-defined]
+            with (
+                device_autocast_ctx,   # type: ignore[attr-defined]
+                torch.amp.autocast("cpu", **cpu_autocast_kwargs),
+                _apply_torch_function_mode_stack(torch_function_mode_stack),
+                recompute_context,
+            ):
                 fn(*args, **kwargs)
 
     new_frame = _CheckpointFrame(