Repository: Lightning-AI/lightning-thunder
Pull request #1956: autodiff processor
Author: login: t-vi
Created: 2025-04-10T21:49:17Z | Merged: 2025-05-20T13:28:31Z

Will clean up add comments etc. to get a discussion basis.

Goals
- model after transform for execution, in particular use TraceSubstitutionProcessor and to the decomposition the same way (by having a "queue" of  bound symbols to process),
- no use of "thunder.common.trace", use push scope instead,
- no conversion between functions and traces,
- (builds the joint trace instead to integrate with @beverlylytle 's work) generate the backward immediately with the augmented forward,

- [TODO] want to integrate more into the transform system,
- [TODO] propagate requires_grad in the aug forward pass,
- [TODO, probably some things vanish on their own now] clear terminology. Not calling bound symbols "symbol", no "VJPDual" when there is no duality involved,
- (seems to be better achieved by joint traces) clearer contract between forward and backward (it seems bad that one could not deduplicate saved for backwards),
- TODO/need to verify: better robustness (e.g. there is no error message if you use get_grad on a non-output currently, and that seemed hard to debug),


It still is somewhat tricky because the backward generation does not fit the tooling well, but it seems like a start for the decomposition and dropping `thunder.common.trace`.

I didn't fix rematerialize_forward_and_backward yet. My understanding is that @beverlylytle 's work will obsolete that (please correct me if I'm wrong.).



Diff:
diff --git a/thunder/core/symbol.py b/thunder/core/symbol.py
index 3ab35e57fd..b35f408fc6 100644
--- a/thunder/core/symbol.py
+++ b/thunder/core/symbol.py
@@ -321,9 +321,8 @@ def __call__(self, *args, **kwargs):
             # vjp transform (applied later).
             def tag_tensorproxy_output_as_detached(proxy):
                 if isinstance(proxy, TensorProxy):
-                    # We need to remove name from trace, otherwise replace will return a proxy with new name.
-                    trace.names.remove(proxy.name)
-                    return proxy.replace(tags=(ProxyTag.DETACHED_AUTOGRAD_GRAPH,))
+                    proxy.tags.add(ProxyTag.DETACHED_AUTOGRAD_GRAPH)
+
                 return proxy
 
             result = tree_map(tag_tensorproxy_output_as_detached, result)
diff --git a/thunder/executors/torch_autograd.py b/thunder/executors/torch_autograd.py
index ce10977181..1a591d1430 100644
--- a/thunder/executors/torch_autograd.py
+++ b/thunder/executors/torch_autograd.py
@@ -225,7 +225,7 @@ def connect_to_autograd(
 
 def split_forward_backward(computation_trc: TraceCtx, compile_data, compile_stats, /, *flat_args):
     from thunder.core.rematerialization import rematerialize_all_gather, rematerialize_forward_and_backward
-    from thunder.core.transforms import forward_and_backward_from_trace
+    from thunder.transforms.autodiff import forward_and_backward_from_trace
     from thunder.distributed.transforms import FSDPCommBucketing
     from thunder.distributed.utils import sort_data_parallel_syncs, sort_waits, sort_communication_ops
     from thunder.executors.passes import del_last_used, transform_for_execution
@@ -432,7 +432,8 @@ def split_forward_backward(computation_trc: TraceCtx, compile_data, compile_stat
     bw_extrace = del_last_used(bw_extrace, clear_mutable_collections=True)
     bw_traces.append(bw_extrace)
 
-    bw_trace = rename_bwd_trace_outputs(bw_extrace, fw_extrace)
+    bw_extrace = rename_bwd_trace_outputs(bw_extrace, fw_extrace)
+    bw_traces.append(bw_extrace)
 
     if compile_stats is not None:
         compile_stats.last_traces += fw_traces
diff --git a/thunder/tests/distributed/test_fsdp.py b/thunder/tests/distributed/test_fsdp.py
index 785df894b2..3963432c26 100644
--- a/thunder/tests/distributed/test_fsdp.py
+++ b/thunder/tests/distributed/test_fsdp.py
@@ -117,7 +117,9 @@ def test_rematerialize_all_gather(self):
         m = ToyModel().to(device)
         cm = thunder.jit(m)
         cm = fsdp(cm, device=device, broadcast_from=0)
-        x = torch.ones((2, 12), device=device)
+        # if we don't require grad of x, we do not need the t_net1_weight
+        # in the backward below and it could be legitimately DCE'd.
+        x = torch.ones((2, 12), device=device, requires_grad=True)
         cm(x).mean().backward()
 
         fwd_trc = [t for t in thunder.last_traces(cm) if thunder.core.trace.TraceTag.AUGMENTED_FORWARD in t.tags][0]
@@ -131,8 +133,8 @@ def test_rematerialize_all_gather(self):
         #       in the original trace and are inputs to all_gather, the unshard are the outputs fo the corresponding wait
         #       If you fix this to be dynamically discerned, you'll be my hero.
         sharded_param_names = ("t_net1_weight", "t_net2_weight")
-        # t20 and t22 are all-gather'ed t_net1_weight and t_net2_weight, respectively.
-        unshard_param_names = ("t20", "t22")
+        # t30 and t84 are all-gather'ed t_net1_weight and t_net2_weight, respectively.
+        unshard_param_names = ("t30", "t84")
         result_saved_for_bwd = [x.name for x in fwd_trc.bound_symbols[-1].args[1][0]]
         self.assertTrue(all(t not in sharded_param_names for t in result_saved_for_bwd))
         self.assertTrue(all(t in result_saved_for_bwd for t in unshard_param_names))
diff --git a/thunder/tests/test_examine_memory.py b/thunder/tests/test_examine_memory.py
index 51278003de..61e2ef086b 100644
--- a/thunder/tests/test_examine_memory.py
+++ b/thunder/tests/test_examine_memory.py
@@ -116,4 +116,4 @@ def test_nanogpt_block():
     assert max_mem_fw[0] == 381754368
     assert sum(max_mem_fw[1].values()) == 375462912
     assert max_mem_bw[0] == 437292032
-    assert sum(max_mem_bw[1].values()) == 40934400
+    assert sum(max_mem_bw[1].values()) == 34642944
diff --git a/thunder/tests/test_grad.py b/thunder/tests/test_grad.py
index fa8bac4b52..6e3d2f854f 100644
--- a/thunder/tests/test_grad.py
+++ b/thunder/tests/test_grad.py
@@ -1856,7 +1856,7 @@ def f(x):
 
     with pytest.raises(
         RuntimeError,
-        match="number of outputs of the original forward function must be the same as the number of primal outputs",
+        match="The number of outputs of the gradient transform function must be the same as the number of outputs of the original forward function.",
     ):
         _ = jf(a)
 
@@ -1944,35 +1944,6 @@ def func(x):
     torch.testing.assert_close(actual_gr, expected_gr)
 
 
-@pytest.mark.parametrize("device", ("cuda", "cpu"))
-def test_backward_recomputation_decomposed_ops(device):
-    if device == "cuda" and not torch.cuda.is_available():
-        pytest.skip("CUDA is not available")
-
-    def fn(a):
-        return torch.nn.functional.gelu(a)
-
-    # rematerialization will also trigger recomputation here.
-    jfn = thunder.jit(fn, executors=())
-    jfn2 = thunder.jit(fn, auto_recompute_intermediates=True)
-    a = torch.randn(2, 2, device=device, requires_grad=True)
-    res = jfn(a)
-    res2 = jfn2(a)
-    assert len(res.grad_fn.saved_tensors) == 3  # should be decomposed
-    assert len(res2.grad_fn.saved_tensors) == 1
-
-    if NVFUSER_AVAILABLE and device == "cuda":
-        # check everything is fused
-        assert {bsym.sym.name for bsym in thunder.last_backward_traces(jfn2)[-1].bound_symbols} == {
-            "nvFusion0",
-            "clear_mutable_collection",
-            "python_return",
-            "python_del",
-            "unpack_sequence",
-            "unpack_trivial",
-        }
-
-
 @requiresCUDA
 def test_benchmark_grad():
     from thunder.benchmarks.utils import backward_only
diff --git a/thunder/transforms/autodiff.py b/thunder/transforms/autodiff.py
new file mode 100644
index 0000000000..3d17770bd3
--- /dev/null
+++ b/thunder/transforms/autodiff.py
@@ -0,0 +1,481 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims, utils
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+    recompute_saved_for_backward,
+)
+from thunder.core.proxies import ProxyTag
+from thunder.core.symbol import BoundSymbol
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+def _should_recompute_bsym_in_backward(bsym):
+    return any((ProxyTag.RECOMPUTE_IN_BACKWARD in o.tags) for o in bsym.flat_proxy_outs)
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backward registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                ########### BEGINNING of return handling (and putting the backward computation in the joint trace)
+                # This is big (and a bit messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in torch.autograd.Function)
+                # The backward computation in self.collected_bw_part_syms is in the order of the
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # which names we saw get_grad for, we cannot do put_grad on them anymore
+                names_seen_get_grad = set()
+
+                # we iterate through the collected symbols backward in order (from the generation of the list, this
+                # is the backward of the last forward instruction first, then then the backward of the second to last etc.
+                for bw_bsyms in self.collected_bw_part_bsyms:
+                    # if we don't have and grad_outs, we know that the grad_ins are not needed, so we don't compute them
+                    # (e.g. for `where(a * b > 0, a, b)` we do not need the grad components of a * b propagated
+                    if any(
+                        ((gg.name in grad_proxy_map) or (gg.name in output_proxy_names))
+                        for gg in (bs.args[0] for bs in bw_bsyms if bs.sym == prims.get_grad)
+                    ):
+                        for nbsym in bw_bsyms:
+                            if nbsym.sym == prims.put_grad:
+                                # put_grad adds something to the gradients
+                                # - if it is an output of the function, we also need to get the
+                                #   "grad_output" (but only once)
+                                # - if we already had a gradient, we need to add the new component
+                                # - structurally (for correctness), there should be no put_grad
+                                #   after a get_grad
+                                p = self.env.get(nbsym.args[0].name, nbsym.args[0])
+                                assert p.name not in names_seen_get_grad, f"put_grad after get_grad on {p.name}"
+                                current_grad = grad_proxy_map.get(p.name)
+                                new_grad = nbsym.args[1]
+                                if current_grad is None and p.name in output_proxy_names:
+                                    self.new_trace.push_scope([])
+                                    current_grad = prims.get_grad(p)
+                                    self.add_processed_bsyms(self.new_trace.pop_scope())
+
+                                if current_grad is not None:
+                                    new_grad = self.add_bsyms_from_function(thunder.torch.add, current_grad, new_grad)
+
+                                grad_proxy_map[p.name] = new_grad
+                                self.write(new_grad, new_grad)
+                            elif nbsym.sym == prims.get_grad:
+                                # get grad reads a gradient, this could be
+                                # - the gradient of an intermediate (possibly combined with an output) form put_grad above
+                                #   note that all put_grads are happen before all get_grads from the dataflow of the forward,
+                                #   in this case we use the gradient from the put_grad(s)
+                                # - an output of the forward function (that has not been combined with an intermediate in put grad),
+                                #   then the gradient is like the "grad_out" in the parameter list of a autograd.Function.backward
+                                #   in this case we keep a get_grad in the trace.
+                                # - an intermediate that does not have a gradient computed. In this case the gradient is 0 in the shape
+                                #   of the get_grad argument. This happens e.g. if we have something like "a > 0" for a float a then
+                                #   a has grad 0.
+                                #   TODO: We could (and should) optimize this (see below).
+                                p = nbsym.args[0]
+
+                                # TODO: set proxy requires grad here!
+                                name = p.name
+                                names_seen_get_grad.add(name)
+                                current_grad = grad_proxy_map.get(name)
+                                if current_grad is not None:
+                                    # do we also need to map?
+                                    self.write(nbsym.output, current_grad)
+                                    # replace_output_with_current_grad
+                                    self.swap_map[thunder.core.proxies.variableify(nbsym.output)] = current_grad
+                                elif name in output_proxy_names:
+                                    # output here???
+                                    new_bsym = nbsym.from_bsym()
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                else:
+                                    # TODO: mark this? if all inputs to the backward formula are unused, we would not want to compute it.
+                                    new_bsym = thunder.torch.zeros.bind(
+                                        *p.shape, device=p.device, dtype=p.dtype, output=nbsym.output
+                                    )
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                    self.write(nbsym.output, new_bsym.output)
+                            else:
+                                # all other symbols are just computation
+                                self.add_processed_bsyms([nbsym.from_bsym()])
+                self.collected_bw_part_bsyms.clear()
+
+                # we collect the gradients fo the flat args in the output dictionary's 'grad_flat_args'
+                # this means that the joint trace has them as outputs, which is good.
+                grad_flat_args = []
+                for p in bsym.args[0]["flat_args"]:
+                    # or p = self.read(p) here?
+                    if isinstance(p, thunder.TensorProxy) and p.requires_grad and p.name in grad_proxy_map:
+                        # is it always OK if we don't have a gradient? (one case: unused input)
+                        # result of put_grad???
+                        grad_flat_args.append(grad_proxy_map[p.name])
+                    else:
+                        grad_flat_args.append(None)
+
+                bsym.args[0]["grad_flat_args"] = grad_flat_args
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+                ########### END of return handling (and putting the backward computation in the joint trace)
+
+            # we now handle various bound symbols by collecting augmented forward and backward symbols
+
+            # 1. constant for gradients (no grad required)
+            #    executing synchronize here terrible hack to cope with non-grad-needing sharded tensors
+            #    as required by LoRA. We should have a symbol tag "always compute grad" instead.
+            if is_constant_for_vjp(bsym) and not bsym.sym.name == "synchronize":
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+
+            # 2. Special case the thunder.torch.checkpoint higher order function
+            if bsym.sym == thunder.torch.checkpoint:
+                # Tag all intermediate outputs as to be recomputed.
+                function_arg_names = {a.name for a in bsym.flat_proxy_args}
+
+                for subsym in bsym.subsymbols:
+                    for o in subsym.flat_proxy_outs:
+                        if o.name not in function_arg_names:
+                            o.tags.add(ProxyTag.RECOMPUTE_IN_BACKWARD)
+
+                # decompose
+                self.add_unprocessed_bsyms(bsym.subsymbols[:])
+                return
+
+            # 3a. see if we have a grad_transform (e.g. from OperatorExecutor.register_grad_transform)
+
+            # executor or global grad transform
+            joint_forward_backward, _ = _get_gradfn_and_executor(bsym)
+
+            if joint_forward_backward is None:
+                # 3b. if we don't have a grad_transform for a bsym, maybe we have old style augmented forward and backward
+                # registered. In this case, we make the joint_forward_backward which works like a grad_gransform.
+                # registered augmented forward impl (where aug fwd and backward are registered separately)
+                aug_fwd_impl = augmented_forward_impls.get(bsym.sym.id)
+                if aug_fwd_impl is not None:
+                    bwd_impl = backward_impls.get(bsym.sym.id)
+
+                    # this is our ad hoc combined forward and backward function ("grad_transform")
+                    def joint_forward_backward(*args, **kwargs):
+                        arg_proxy_names = {a.name for a in bsym.flat_proxy_args}
+
+                        # run the augmented forward res
+                        aug_fwd_res = aug_fwd_impl(*bsym.args, **bsym.kwargs)
+
+                        # aug_fwd_res could be either VJPDual or just a tuple, in any case, we can decompose it.
+                        res, saved_for_backward = aug_fwd_res
+
+                        # we need to shallow copy inputs that are returned for "get_grad" and "put_grad" to properly work
+                        # (this shallow copy is the equivalent because we of an "edge" in the PyTorch autograd graph)
+                        def shallow_copy_if_input(p):
+                            if isinstance(p, thunder.TensorProxy) and p.name in arg_proxy_names:
+                                return thunder.core.prims.shallow_copy(p)
+                            return p
+
+                        res = tree_map(shallow_copy_if_input, res)
+
+                        # now we need the backward. it starts by getting the grad_outs
+                        grad_outs = []
+                        for r in thunder.core.pytree.tree_iter(res):
+                            if isinstance(r, thunder.TensorProxy):
+                                grad_outs.append(prims.get_grad(r))
+
+                        # The backward computes the grad_inps of the bsym from the grad_outs
+                        # TODO: non-grad outputs of bwd?
+                        grad_inps = bwd_impl(*saved_for_backward, *grad_outs)
+                        if isinstance(grad_inps, thunder.Proxy):
+                            grad_inps = [grad_inps]
+
+                        # match the grad_inps to the inputs of the boudnd symbol and put the grads
+
+                        flat_inps = args
+                        # for autograd_function_apply, skip the function args
+                        # TODO: fix the returned gradients to include two None?.
+                        if bsym.sym == thunder.torch.autograd_function_apply:
+                            flat_inps = args[2:]
+
+                        # there may be non-gradient requiring additional args (todo: maybe only support this for non-tensor ones?)
+                        num_flat_tensor_inps = sum(isinstance(i, thunder.TensorProxy) for i in flat_inps)
+                        utils.check(
+                            num_flat_tensor_inps <= len(grad_inps),
+                            lambda: f"Backward for {bsym.sym.id} returned {len(grad_inps)} value(s), but expected {num_flat_tensor_inps}",
+                        )
+
+                        assert len(grad_inps) <= len(flat_inps)
+                        for i, gi in zip(flat_inps, grad_inps):
+                            # for integer proxies etc. we expect gi to be None
+                            if isinstance(i, thunder.TensorProxy) and gi is not None:
+                                prims.put_grad(i, gi)
+                        return res
+
+            # 3c if we have a joint forward backward (either grad_transform or the ad hoc constructed one above), we can comptue the gradient
+            if joint_forward_backward is not None:
+                self.new_trace.push_scope([])
+                result = joint_forward_backward(*bsym.args, **bsym.kwargs)
+
+                # Check that the number of outputs of the original forward function is the
+                # same as the number of primal outputs of the augmented forward trace
+                utils.check(
+                    len(utils.sequencify(bsym.output)) == len(utils.sequencify(result)),
+                    lambda: f"While generating forward and backward functions for {bsym.sym.name}, encountered an error.\n"
+                    "The number of outputs of the gradient transform function must be the same as the number of outputs of the original forward function.\n"
+                    f"Number of outputs of the original forward function: {len(utils.sequencify(bsym.output))}\n"
+                    f"Number of primal outputs of the gradient transform / augmented forward: {len(utils.sequencify(result))}\n"
+                    "Please check the forward function and the gradient transform function / augmented forward to ensure that they have the same number of outputs.",
+                )
+                self.set_result(result)
+                new_bsyms = self.new_trace.pop_scope()
+
+                # simple splitting: only compute in forward what is needed for the output
+                forward_part_proxy_names = {
+                    o.name for o in thunder.core.pytree.tree_iter(result) if isinstance(o, thunder.Proxy)
+                }
+                forward_part_bsyms = []
+                backward_part_bsyms = []
+                for nbsym in reversed(new_bsyms):
+                    # argnames = {p.name for p in nbsym.flat_proxy_args}
+                    if nbsym.sym in {prims.get_grad, prims.put_grad}:
+                        backward_part_bsyms.insert(0, nbsym)
+                        continue
+                    assert nbsym.sym != prims.unpack_trivial  # can unpack trivial happen here?
+                    if any((o.name in forward_part_proxy_names) for o in nbsym.flat_proxy_outs):
+                        forward_part_bsyms.insert(0, nbsym)
+                        bsym_arg_names = {a.name for a in nbsym.flat_proxy_args}
+                        bsym_out_names = {o.name for o in nbsym.flat_proxy_outs if o.name not in bsym_arg_names}
+                        forward_part_proxy_names.update(bsym_arg_names)
+                        forward_part_proxy_names.update(bsym_out_names)
+                        continue
+                    backward_part_bsyms.insert(0, nbsym)
+                self.add_processed_bsyms(forward_part_bsyms)
+
+                self.collected_bw_part_bsyms.insert(0, backward_part_bsyms)
+                return
+
+            # 4. No gradient transform found, need to decompose the symbol
+
+            # error if this is a primitive
+            utils.check(not bsym.sym.is_prim, lambda: f"Failed to find a gradient transform for bound symbol {bsym=}")
+
+            # TODO: check if this is needed: the old impl checked whether len(bsym.subsymbols) > 0 except for the special case "torch.nn.functional.dropout" with p=0...
+            # add the decomposition (= the subsymbols) to the front of the symbols to be processed
+            self.add_unprocessed_bsyms(bsym.subsymbols[:])
+
+            # end of 4 and end of the bsym processing loop.
+
+    start_time_ns = time.perf_counter_ns()
+
+    # run the trace through the processor
+    trace, _ = AugmentedForwardProcessor(trace)()
+    # run through DCE in case some of the gradients of intermediates are not needed.
+    trace = thunder.core.transform_common.dce(trace)
+
+    end_time_ns = time.perf_counter_ns()
+    elapsed_time_ns = end_time_ns - start_time_ns
+    elapsed_time_millis = elapsed_time_ns // 1000000
+    trace.set_provenance(
+        thunder.core.trace.TraceProvenance(f"Grad transform pass (took {elapsed_time_millis} milliseconds)")
+    )
+    return trace
+
+
+def split_into_forward_and_backward(joint_trace):
+    """split a joint trace for forward and backward into separate ones, including recomputation (aka activation checkpointing)"""
+
+    # the joint trace will have the forward computation at the beginning and then the backward computation
+    # from how it is constructed.
+    # we split the trace:
+    # - forward symbols go into forward_part_bsyms
+    # - all symbols not in the forward go into backward_part_bsyms
+    # - for recomputation (aka activation checkpointing), we want to insert symbols going into the forward also into the
+    #   backward, but we want to do so "just in time". To this end, we gather the symbols in a dict and later
+    #   insert it when their respective outputs are needed. This is in backward_part_bsyms_recomputed
+    #   The just in time recomputation is a heuristic to save memory mimicking checkpointing: e.g. for a checkpointed
+    #   block, the forward would be recomputed just before computing the gradient.
+    # the splitting is done in the reverse order of the bound symbols, and works out which bits are needed for the forward
+    # from there.
+    forward_part_bsyms = []
+    backward_part_bsyms = []
+    backward_part_bsyms_recomputed: dict[str, tuple[BoundSymbol, set[str]]] = {}  # name -> (bsym, name_set)
+
+    # return has a dict, it contains the forward outputs in "output" and the backward outputs in "grad_flat_args", corresponding
+    # to "flat_args"
+    return_bsym = joint_trace.bound_symbols[-1]
+    assert return_bsym.sym == prims.python_return
+    fw_output = return_bsym.args[0]["output"]
+    assert isinstance(fw_output, tuple)
+
+    grad_outs = [None for _ in fw_output]
+    output_pos = {o.name: i for i, o in enumerate(fw_output) if isinstance(o, thunder.TensorProxy)}
+
+    # the proxies we need to compute in the forward - we start with the outputs of the forward
+    forward_proxy_names = {o.name for o in thunder.core.pytree.tree_iter(fw_output) if isinstance(o, thunder.Proxy)}
+    # we also have the inputs available, so we add flat_args.
+    # for inplace, we need to update this (or have flat args be the right thing?...)
+    forward_proxy_names.update(a.name for a in return_bsym.args[0]["flat_args"] if isinstance(a, thunder.Proxy))
+
+    # We keep track of the names of proxies we recompute in the backward as those will not need to be part of the
+    # ones saved in the forward for the backward
+    backward_recomputed_proxy_names = set()
+
+    # loop over the bound symbols in reverse (see above)
+    for bsym in reversed(joint_trace.bound_symbols):
+        if bsym.sym == prims.python_return:  # we will re-do the return statements below, so we skip it here
+            continue
+
+        # unpack_trivial is always a (forward trace) argument, the backward inputs are from get_grad
+        # anything that outputs something needed for the forward (the forward_proxy_names) is part of the forward
+        if any((o.name in forward_proxy_names) for o in bsym.flat_proxy_outs) or bsym.sym == prims.unpack_trivial:
+            forward_part_bsyms.insert(0, bsym.from_bsym())
+            bsym_arg_names = {a.name for a in bsym.flat_proxy_args}
+            bsym_out_names = {o.name for o in bsym.flat_proxy_outs if o.name not in bsym_arg_names}
+            forward_proxy_names.update(bsym_arg_names)
+            forward_proxy_names.update(bsym_out_names)
+
+            # if we need to recompute this, we also add the bound symbol to backward_part_bsyms_recomputed
+            if _should_recompute_bsym_in_backward(bsym):
+                backward_recomputed_proxy_names.update(bsym_out_names)
+                bsym_rec = (bsym.from_bsym(), bsym_out_names)
+                backward_part_bsyms_recomputed.update((n, bsym_rec) for n in bsym_out_names)
+
+            continue
+
+        # get grad is always part of the input, record the grad_out (will be part of the "cotangents" list)
+        if bsym.sym == prims.get_grad:
+            grad_outs[output_pos[bsym.args[0].name]] = bsym.output
+            continue
+
+        # copy_ updating a forward proxy is special regardless of the output
+        if bsym.sym == prims.copy_ and bsym.args[1].name in forward_proxy_names:
+            # todo: should we also handle ltorch.copy_ ?
+            forward_part_bsyms.insert(0, bsym.from_bsym())
+            forward_proxy_names.update(a.name for a in bsym.flat_proxy_args)
+            continue
+
+        # if we don't need to have it in the forward, it is part of the backward
+        backward_part_bsyms.insert(0, bsym.from_bsym())
+
+    # we insert the recomputed symbols just before where they are needed.
+    # as this inserts into the list during processing, we use a while loop rather than
+    # a for loop
+    bw_idx = 0
+    while bw_idx < len(backward_part_bsyms):
+        bsym = backward_part_bsyms[bw_idx]
+        modified = False
+        for n in reversed(bsym.flat_proxy_args):
+            recomp_bsym_rec = backward_part_bsyms_recomputed.get(n.name)
+            if recomp_bsym_rec is not None:
+                modified = True
+                recomp_bsym, recomp_output = recomp_bsym_rec
+                backward_part_bsyms.insert(bw_idx, recomp_bsym)
+                for nn in recomp_output:
+                    del backward_part_bsyms_recomputed[nn]
+        if not modified:
+            bw_idx += 1
+
+    # collect proxies needing to be saved in the forward for the backward
+    saved_for_backward = {}
+    for bsym in backward_part_bsyms:
+        saved_for_backward.update(
+            (a.name, a)
+            for a in bsym.flat_proxy_args
+            if a.name in forward_proxy_names and a.name not in backward_recomputed_proxy_names
+        )
+    saved_for_backward_tensors = [p for p in saved_for_backward.values() if isinstance(p, thunder.TensorProxy)]
+    saved_for_backward_other = [p for p in saved_for_backward.values() if not isinstance(p, thunder.TensorProxy)]
+
+    # we build the forward trace
+    forward_trace = thunder.core.trace.from_trace(joint_trace)
+    forward_trace.tags.add(thunder.core.trace.TraceTag.AUGMENTED_FORWARD)
+    forward_trace.names = forward_trace.names.copy()  ## ehem
+    forward_trace.bound_symbols += forward_part_bsyms
+
+    # now we create the return value and return bound symbol for the forward
+    fw_output_dict = {k: v for k, v in return_bsym.args[0].items() if k != "grad_flat_args"}
+    flat_output, _ = thunder.core.pytree.tree_flatten_with_dataclass(fw_output)
+    fw_output_dict["flat_output"] = tuple(flat_output)
+    with thunder.core.trace.tracectx(forward_trace):
+        prims.python_return(fw_output_dict, (saved_for_backward_tensors, saved_for_backward_other))
+
+    # then we construct the backward trace, unpacking saved_for_backward and cotangents lists
+    def backward_fn(saved_for_backward, cotangents):
+        pass
+
+    backward_trace = thunder.core.trace.TraceCtx(fn=backward_fn)
+    backward_trace.names = forward_trace.names
+    backward_trace.name_ctr = forward_trace.name_ctr
+
+    # TODO: make this name-agnostic
+    backward_trace.names.discard("saved_for_backward")
+    backward_trace.names.discard("cotangents")
+
+    # set up the inputs of the backward properly (args and unpacking)
+    with thunder.core.trace.tracectx(backward_trace):
+        p_C0 = thunder.core.proxies.CollectionProxy(list(saved_for_backward_tensors), name="C0")
+        p_C1 = thunder.core.proxies.CollectionProxy(list(saved_for_backward_other), name="C1")
+        p_saved_for_backward = thunder.core.proxies.CollectionProxy([p_C0, p_C1], name="saved_for_backward")
+        p_cotangents = thunder.core.proxies.CollectionProxy(grad_outs, name="cotangents")
+
+        # set the args (which currently don't use the collection proxies but the collections directly)
+        saved_for_backward_tuple = [p_C0.collection(), p_C1.collection()]
+        backward_trace.args = (saved_for_backward_tuple, p_cotangents.collection())
+
+        # unpacking symbols
+        prims.unpack_trivial(p_saved_for_backward, name="saved_for_backward")
+        prims.unpack_trivial(p_cotangents, name="cotangents")
+        prims.unpack_sequence(p_saved_for_backward, len(p_saved_for_backward.coll))
+        prims.unpack_sequence(p_cotangents, len(p_cotangents.coll))
+        prims.unpack_sequence(p_C0, len(p_C0.coll))
+        prims.unpack_sequence(p_C1, len(p_C1.coll))
+
+    # add the backward computation
+    backward_trace.bound_symbols += backward_part_bsyms
+
+    # and finally the backward return statement
+    with thunder.core.trace.tracectx(backward_trace):
+        prims.python_return(tuple(return_bsym.args[0]["grad_flat_args"]))
+
+    return forward_trace, backward_trace
+
+
+def forward_and_backward_from_trace(trace: thunder.core.trace.TraceCtx, torch_autograd=False) -> ForwardBackwardTraces:
+    if not torch_autograd:
+        return thunder.core.transforms.forward_and_backward_from_trace(trace, torch_autograd=torch_autograd)
+
+    joint_trace = grad_transform_on_trace(trace)
+
+    forward_trace, backward_trace = split_into_forward_and_backward(joint_trace)
+    return ForwardBackwardTraces(forward_trace, backward_trace)


Reviews:
user: login: IvanYashchuk
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-11T10:01:52Z
user: login: t-vi
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-13T15:49:11Z
user: login: t-vi
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-13T15:50:25Z
user: login: beverlylytle
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-15T08:17:27Z
user: login: t-vi
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-15T11:17:14Z
user: login: t-vi
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-15T11:22:06Z
user: login: t-vi
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-15T15:53:05Z
user: login: beverlylytle
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-16T16:06:38Z
user: login: t-vi
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T14:44:23Z
user: login: t-vi
body: There are two details here in the surroundings of the autodiff itself.
The other changes outside the thunder.transforms.autodiff are benign updates of tests.

state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-19T19:53:59Z
user: login: t-vi
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-19T19:54:19Z
user: login: lantiga
body: Epic 🚀, just a few nits here and there
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T12:11:20Z
user: login: beverlylytle
body: Looks good!  Only some nits.
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-20T12:30:06Z
user: login: t-vi
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-20T12:51:05Z
user: login: t-vi
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-20T12:55:11Z

Review comments:
diff_hunk: @@ -0,0 +1,281 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+)
+from thunder.core.utils import check
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
path: thunder/transforms/autodiff.py
user: login: IvanYashchuk
body: That's cool that it works, but these more than 100 lines of code can be condensed into <10 lines of code. In the current prototype, every operation in the input trace is reevaluated/reinterpreted, redirecting all calls to the augmented forward and backward functions. Here's the condensed version of the same logic:
```py
from thunder.prims import get_grad
from thunder.core.transforms import put_grads, augmented_forward_pass, backward_pass

def grad_transform_on_trace(trace, /, *args, **kwargs):
    def joint_forward_and_backward(*args, **kwargs):
        result, all_intermediates = augmented_forward_pass(*args, trace=trace, **kwargs)
        output_grad = get_grad(result)
        backward_result = backward_pass(all_intermediates, trace, cotangents)
        return {"output": result, "grad_flat_args": backward_result}
    return trace(use_dce=False)(joint_forward_and_backward, *args, **kwargs)
```
created_at: 2025-04-11T10:01:52Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,281 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+)
+from thunder.core.utils import check
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
path: thunder/transforms/autodiff.py
user: login: t-vi
body: Yeah, but how many LOC is augmented forward pass and backward_pass? How many LOC is trace?
All these have to go, and this is showing them out.
created_at: 2025-04-13T15:49:11Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,281 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+)
+from thunder.core.utils import check
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
path: thunder/transforms/autodiff.py
user: login: t-vi
body: Not to forget that there is an entire caching mechanism we try to drop here.
created_at: 2025-04-13T15:50:25Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,301 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+)
+from thunder.core.utils import check
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                # This is big (and a big messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols in order
+                for nbsym in self.collected_bw_part_bsyms:
+                    if nbsym.sym == prims.put_grad:
+                        # put_grad adds something to the gradients
+                        # - if it is an output of the function, we also need to get the
+                        #   "grad_output" (but only once)
+                        # - if we already had a gradient, we need to add the new component
+                        # - structurally (for correctness), there should be no put_grad
+                        #   after a get_grad
+                        p = self.env.get(nbsym.args[0].name, nbsym.args[0])
+                        current_grad = grad_proxy_map.get(p.name)
+                        new_grad = nbsym.args[1]
+                        if current_grad is None and p.name in output_proxy_names:
+                            self.new_trace.push_scope([])
+                            current_grad = prims.get_grad(p)
+                            self.add_processed_bsyms(self.new_trace.pop_scope())
+
+                        if current_grad is not None:
+                            new_grad = self.add_bsyms_from_function(thunder.torch.add, current_grad, new_grad)
+
+                        grad_proxy_map[p.name] = new_grad
+                        self.write(new_grad, new_grad)
+                    elif nbsym.sym == prims.get_grad:
+                        p = nbsym.args[0]
+                        # TODO: set requires grad here!
+                        name = p.name
+                        current_grad = grad_proxy_map.get(name)
+                        if current_grad is not None:
+                            # do we also need to map?
+                            self.write(nbsym.output, current_grad)
+                            self.swap_map[thunder.core.proxies.variableify(nbsym.output)] = current_grad
+                            # XXX_replace_output_with_current_grad
+                        elif name in output_proxy_names:
+                            new_grad = self.add_processed_bsyms([nbsym.from_bsym()])
+                            grad_proxy_map[name] = new_grad
+                        else:
+                            raise RuntimeError(f"grad of non-output and non-intermediate {name} requested")
+                    else:
+                        self.add_processed_bsyms([nbsym.from_bsym()])
+                self.collected_bw_part_bsyms.clear()
+
+                grad_flat_args = []
+                for p in bsym.args[0]["flat_args"]:
+                    # or p = self.read(p) here?
+                    if isinstance(p, thunder.Proxy) and p.requires_grad and p.name in grad_proxy_map:
+                        # is it always OK if we don't have a gradient? (one case: unused input)
+                        # result of put_grad???
+                        grad_flat_args.append(grad_proxy_map[p.name])
+                    else:
+                        grad_flat_args.append(None)
+
+                bsym.args[0]["grad_flat_args"] = grad_flat_args
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+
+            # constant for gradients (no grad required)
+            if is_constant_for_vjp(bsym):
+                if bsym.sym.name == "synchronize":
+                    # This is a *really* terrible hack to cope with non-grad-needing sharded tensors
+                    # as required by LoRA.
+                    from thunder.distributed.prims import all_gather
+
+                    def synchronize_impl(a, group):
+                        return all_gather(a, group, True).wait()
+
+                    self.add_bsyms_from_function(synchronize_impl, *bsym.args, **bsym.kwargs)
+                    return
+
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+
+            # executor or global grad transform
+            joint_forward_backward, _ = _get_gradfn_and_executor(bsym)
+
+            if joint_forward_backward is None:
+                # registered augmented forward impl (where aug fwd and backward are registered separately)
+                aug_fwd_impl = augmented_forward_impls.get(bsym.sym.id)
+                if aug_fwd_impl is not None:
+                    bwd_impl = backward_impls.get(bsym.sym.id)
+
+                    def joint_forward_backward(*args, **kwargs):
+                        aug_fwd_res = aug_fwd_impl(*bsym.args, **bsym.kwargs)
+                        if isinstance(aug_fwd_res, thunder.core.transforms.VJPDual):
+                            res, saved_for_backward = aug_fwd_res
+                        else:
+                            res, saved_for_backward = aug_fwd_res
+                        grad_outs = []
+                        for r in thunder.core.pytree.tree_iter(res):
+                            if isinstance(r, thunder.TensorProxy):
+                                grad_outs.append(prims.get_grad(r))
+                        # to do non-grad outputs of bwd
+                        grad_inps = bwd_impl(*saved_for_backward, *grad_outs)
+                        if isinstance(grad_inps, thunder.Proxy):
+                            grad_inps = [grad_inps]
+
+                        flat_inps = args
+
+                        for i, gi in zip(flat_inps, grad_inps, strict=True):  # strict?
+                            if isinstance(i, thunder.TensorProxy):
+                                prims.put_grad(i, gi)
+                        return res
+
+            if joint_forward_backward is not None:
+                self.new_trace.push_scope([])
+                result = joint_forward_backward(*bsym.args, **bsym.kwargs)
+                self.set_result(result)
+                new_bsyms = self.new_trace.pop_scope()
+
+                forward_part_bsyms = []
+                backward_part_bsyms = []
+                backward_proxy_names = set()
+                for nbsym in new_bsyms:
+                    argnames = {p.name for p in nbsym.flat_proxy_args}
+                    if nbsym.sym in {prims.get_grad, prims.put_grad} or any(
+                        n in backward_proxy_names for n in argnames
+                    ):
+                        backward_proxy_names.update(p.name for p in nbsym.flat_proxy_outs if p.name not in argnames)
+                        backward_part_bsyms.append(nbsym)
+                    else:
+                        forward_part_bsyms.append(nbsym)
+
+                self.add_processed_bsyms(forward_part_bsyms)
+
+                self.collected_bw_part_bsyms[:0] = backward_part_bsyms
+                return
+
+            # No gradient transform found, need to descend
+            check(not bsym.sym.is_prim, lambda: f"Failed to find a gradient transform for bound symbol {bsym=}")
+
+            # TODO: check if this is needed: the old impl checked whether len(bsym.subsymbols) > 0 except for the special case "torch.nn.functional.dropout" with p=0...
+
+            # OUTPUTS to map
+            self.add_unprocessed_bsyms(bsym.subsymbols[:])
+
+    start_time_ns = time.perf_counter_ns()
+
+    processor = AugmentedForwardProcessor(trace)
+    trace, _ = processor()
+    # trace = thunder.core.transform_common.dce(trace)
+
+    end_time_ns = time.perf_counter_ns()
+    elapsed_time_ns = end_time_ns - start_time_ns
+    elapsed_time_millis = elapsed_time_ns // 1000000
+    trace.set_provenance(
+        thunder.core.trace.TraceProvenance(f"Grad transform pass (took {elapsed_time_millis} milliseconds)")
+    )
+    return trace
+
+
+def split_into_forward_and_backward(joint_trace):
+    forward_part_bsyms = []
+    backward_part_bsyms = []
+
+    return_bsym = joint_trace.bound_symbols[-1]
+    assert return_bsym.sym == prims.python_return
+
+    fw_output = return_bsym.args[0]["output"]
+    assert isinstance(fw_output, tuple)
+
+    grad_outs = [None for _ in fw_output]
+
+    output_pos = {o.name: i for i, o in enumerate(fw_output) if isinstance(o, thunder.TensorProxy)}
+    forward_proxy_names = {o.name for o in thunder.core.pytree.tree_iter(fw_output) if isinstance(o, thunder.Proxy)}
+
+    for bsym in reversed(joint_trace.bound_symbols):
+        if bsym.sym == prims.python_return:
+            continue
+
+        if any((o.name in forward_proxy_names) for o in bsym.flat_proxy_outs):
+            forward_part_bsyms.insert(0, bsym.from_bsym())
+            forward_proxy_names.update(a.name for a in bsym.flat_proxy_args)
+            continue
+
+        if bsym.sym == prims.get_grad:
+            grad_outs[output_pos[bsym.args[0].name]] = bsym.output
+            continue
+
+        backward_part_bsyms.insert(0, bsym.from_bsym())
+
+    # collect needed computation
+    saved_for_backward = {}
+    for bsym in backward_part_bsyms:
+        saved_for_backward.update(
+            (a.name, a)
+            for a in bsym.flat_proxy_args
+            if a.name in forward_proxy_names and a.name not in backward_part_bsyms
path: thunder/transforms/autodiff.py
user: login: beverlylytle
body: `a.name not in backward_part_bsyms` is always true. Did you mean something else here?
created_at: 2025-04-15T08:17:27Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,301 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+)
+from thunder.core.utils import check
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                # This is big (and a big messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols in order
+                for nbsym in self.collected_bw_part_bsyms:
+                    if nbsym.sym == prims.put_grad:
+                        # put_grad adds something to the gradients
+                        # - if it is an output of the function, we also need to get the
+                        #   "grad_output" (but only once)
+                        # - if we already had a gradient, we need to add the new component
+                        # - structurally (for correctness), there should be no put_grad
+                        #   after a get_grad
+                        p = self.env.get(nbsym.args[0].name, nbsym.args[0])
+                        current_grad = grad_proxy_map.get(p.name)
+                        new_grad = nbsym.args[1]
+                        if current_grad is None and p.name in output_proxy_names:
+                            self.new_trace.push_scope([])
+                            current_grad = prims.get_grad(p)
+                            self.add_processed_bsyms(self.new_trace.pop_scope())
+
+                        if current_grad is not None:
+                            new_grad = self.add_bsyms_from_function(thunder.torch.add, current_grad, new_grad)
+
+                        grad_proxy_map[p.name] = new_grad
+                        self.write(new_grad, new_grad)
+                    elif nbsym.sym == prims.get_grad:
+                        p = nbsym.args[0]
+                        # TODO: set requires grad here!
+                        name = p.name
+                        current_grad = grad_proxy_map.get(name)
+                        if current_grad is not None:
+                            # do we also need to map?
+                            self.write(nbsym.output, current_grad)
+                            self.swap_map[thunder.core.proxies.variableify(nbsym.output)] = current_grad
+                            # XXX_replace_output_with_current_grad
+                        elif name in output_proxy_names:
+                            new_grad = self.add_processed_bsyms([nbsym.from_bsym()])
+                            grad_proxy_map[name] = new_grad
+                        else:
+                            raise RuntimeError(f"grad of non-output and non-intermediate {name} requested")
+                    else:
+                        self.add_processed_bsyms([nbsym.from_bsym()])
+                self.collected_bw_part_bsyms.clear()
+
+                grad_flat_args = []
+                for p in bsym.args[0]["flat_args"]:
+                    # or p = self.read(p) here?
+                    if isinstance(p, thunder.Proxy) and p.requires_grad and p.name in grad_proxy_map:
+                        # is it always OK if we don't have a gradient? (one case: unused input)
+                        # result of put_grad???
+                        grad_flat_args.append(grad_proxy_map[p.name])
+                    else:
+                        grad_flat_args.append(None)
+
+                bsym.args[0]["grad_flat_args"] = grad_flat_args
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+
+            # constant for gradients (no grad required)
+            if is_constant_for_vjp(bsym):
+                if bsym.sym.name == "synchronize":
+                    # This is a *really* terrible hack to cope with non-grad-needing sharded tensors
+                    # as required by LoRA.
+                    from thunder.distributed.prims import all_gather
+
+                    def synchronize_impl(a, group):
+                        return all_gather(a, group, True).wait()
+
+                    self.add_bsyms_from_function(synchronize_impl, *bsym.args, **bsym.kwargs)
+                    return
+
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+
+            # executor or global grad transform
+            joint_forward_backward, _ = _get_gradfn_and_executor(bsym)
+
+            if joint_forward_backward is None:
+                # registered augmented forward impl (where aug fwd and backward are registered separately)
+                aug_fwd_impl = augmented_forward_impls.get(bsym.sym.id)
+                if aug_fwd_impl is not None:
+                    bwd_impl = backward_impls.get(bsym.sym.id)
+
+                    def joint_forward_backward(*args, **kwargs):
+                        aug_fwd_res = aug_fwd_impl(*bsym.args, **bsym.kwargs)
+                        if isinstance(aug_fwd_res, thunder.core.transforms.VJPDual):
+                            res, saved_for_backward = aug_fwd_res
+                        else:
+                            res, saved_for_backward = aug_fwd_res
+                        grad_outs = []
+                        for r in thunder.core.pytree.tree_iter(res):
+                            if isinstance(r, thunder.TensorProxy):
+                                grad_outs.append(prims.get_grad(r))
+                        # to do non-grad outputs of bwd
+                        grad_inps = bwd_impl(*saved_for_backward, *grad_outs)
+                        if isinstance(grad_inps, thunder.Proxy):
+                            grad_inps = [grad_inps]
+
+                        flat_inps = args
+
+                        for i, gi in zip(flat_inps, grad_inps, strict=True):  # strict?
+                            if isinstance(i, thunder.TensorProxy):
+                                prims.put_grad(i, gi)
+                        return res
+
+            if joint_forward_backward is not None:
+                self.new_trace.push_scope([])
+                result = joint_forward_backward(*bsym.args, **bsym.kwargs)
+                self.set_result(result)
+                new_bsyms = self.new_trace.pop_scope()
+
+                forward_part_bsyms = []
+                backward_part_bsyms = []
+                backward_proxy_names = set()
+                for nbsym in new_bsyms:
+                    argnames = {p.name for p in nbsym.flat_proxy_args}
+                    if nbsym.sym in {prims.get_grad, prims.put_grad} or any(
+                        n in backward_proxy_names for n in argnames
+                    ):
+                        backward_proxy_names.update(p.name for p in nbsym.flat_proxy_outs if p.name not in argnames)
+                        backward_part_bsyms.append(nbsym)
+                    else:
+                        forward_part_bsyms.append(nbsym)
+
+                self.add_processed_bsyms(forward_part_bsyms)
+
+                self.collected_bw_part_bsyms[:0] = backward_part_bsyms
+                return
+
+            # No gradient transform found, need to descend
+            check(not bsym.sym.is_prim, lambda: f"Failed to find a gradient transform for bound symbol {bsym=}")
+
+            # TODO: check if this is needed: the old impl checked whether len(bsym.subsymbols) > 0 except for the special case "torch.nn.functional.dropout" with p=0...
+
+            # OUTPUTS to map
+            self.add_unprocessed_bsyms(bsym.subsymbols[:])
+
+    start_time_ns = time.perf_counter_ns()
+
+    processor = AugmentedForwardProcessor(trace)
+    trace, _ = processor()
+    # trace = thunder.core.transform_common.dce(trace)
+
+    end_time_ns = time.perf_counter_ns()
+    elapsed_time_ns = end_time_ns - start_time_ns
+    elapsed_time_millis = elapsed_time_ns // 1000000
+    trace.set_provenance(
+        thunder.core.trace.TraceProvenance(f"Grad transform pass (took {elapsed_time_millis} milliseconds)")
+    )
+    return trace
+
+
+def split_into_forward_and_backward(joint_trace):
+    forward_part_bsyms = []
+    backward_part_bsyms = []
+
+    return_bsym = joint_trace.bound_symbols[-1]
+    assert return_bsym.sym == prims.python_return
+
+    fw_output = return_bsym.args[0]["output"]
+    assert isinstance(fw_output, tuple)
+
+    grad_outs = [None for _ in fw_output]
+
+    output_pos = {o.name: i for i, o in enumerate(fw_output) if isinstance(o, thunder.TensorProxy)}
+    forward_proxy_names = {o.name for o in thunder.core.pytree.tree_iter(fw_output) if isinstance(o, thunder.Proxy)}
+
+    for bsym in reversed(joint_trace.bound_symbols):
+        if bsym.sym == prims.python_return:
+            continue
+
+        if any((o.name in forward_proxy_names) for o in bsym.flat_proxy_outs):
+            forward_part_bsyms.insert(0, bsym.from_bsym())
+            forward_proxy_names.update(a.name for a in bsym.flat_proxy_args)
+            continue
+
+        if bsym.sym == prims.get_grad:
+            grad_outs[output_pos[bsym.args[0].name]] = bsym.output
+            continue
+
+        backward_part_bsyms.insert(0, bsym.from_bsym())
+
+    # collect needed computation
+    saved_for_backward = {}
+    for bsym in backward_part_bsyms:
+        saved_for_backward.update(
+            (a.name, a)
+            for a in bsym.flat_proxy_args
+            if a.name in forward_proxy_names and a.name not in backward_part_bsyms
path: thunder/transforms/autodiff.py
user: login: t-vi
body: Indeed. Thank you for catching.
created_at: 2025-04-15T11:17:14Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,281 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+)
+from thunder.core.utils import check
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
path: thunder/transforms/autodiff.py
user: login: t-vi
body: But we should absolutely go with the above if it unblocks @beverlylytle 's to make it easier to progress in parallel. Thanks @IvanYashchuk 
created_at: 2025-04-15T11:22:06Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,301 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+)
+from thunder.core.utils import check
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                # This is big (and a big messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols in order
+                for nbsym in self.collected_bw_part_bsyms:
+                    if nbsym.sym == prims.put_grad:
+                        # put_grad adds something to the gradients
+                        # - if it is an output of the function, we also need to get the
+                        #   "grad_output" (but only once)
+                        # - if we already had a gradient, we need to add the new component
+                        # - structurally (for correctness), there should be no put_grad
+                        #   after a get_grad
+                        p = self.env.get(nbsym.args[0].name, nbsym.args[0])
+                        current_grad = grad_proxy_map.get(p.name)
+                        new_grad = nbsym.args[1]
+                        if current_grad is None and p.name in output_proxy_names:
+                            self.new_trace.push_scope([])
+                            current_grad = prims.get_grad(p)
+                            self.add_processed_bsyms(self.new_trace.pop_scope())
+
+                        if current_grad is not None:
+                            new_grad = self.add_bsyms_from_function(thunder.torch.add, current_grad, new_grad)
+
+                        grad_proxy_map[p.name] = new_grad
+                        self.write(new_grad, new_grad)
+                    elif nbsym.sym == prims.get_grad:
+                        p = nbsym.args[0]
+                        # TODO: set requires grad here!
+                        name = p.name
+                        current_grad = grad_proxy_map.get(name)
+                        if current_grad is not None:
+                            # do we also need to map?
+                            self.write(nbsym.output, current_grad)
+                            self.swap_map[thunder.core.proxies.variableify(nbsym.output)] = current_grad
+                            # XXX_replace_output_with_current_grad
+                        elif name in output_proxy_names:
+                            new_grad = self.add_processed_bsyms([nbsym.from_bsym()])
+                            grad_proxy_map[name] = new_grad
+                        else:
+                            raise RuntimeError(f"grad of non-output and non-intermediate {name} requested")
+                    else:
+                        self.add_processed_bsyms([nbsym.from_bsym()])
+                self.collected_bw_part_bsyms.clear()
+
+                grad_flat_args = []
+                for p in bsym.args[0]["flat_args"]:
+                    # or p = self.read(p) here?
+                    if isinstance(p, thunder.Proxy) and p.requires_grad and p.name in grad_proxy_map:
+                        # is it always OK if we don't have a gradient? (one case: unused input)
+                        # result of put_grad???
+                        grad_flat_args.append(grad_proxy_map[p.name])
+                    else:
+                        grad_flat_args.append(None)
+
+                bsym.args[0]["grad_flat_args"] = grad_flat_args
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+
+            # constant for gradients (no grad required)
+            if is_constant_for_vjp(bsym):
+                if bsym.sym.name == "synchronize":
+                    # This is a *really* terrible hack to cope with non-grad-needing sharded tensors
+                    # as required by LoRA.
+                    from thunder.distributed.prims import all_gather
+
+                    def synchronize_impl(a, group):
+                        return all_gather(a, group, True).wait()
+
+                    self.add_bsyms_from_function(synchronize_impl, *bsym.args, **bsym.kwargs)
+                    return
+
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+
+            # executor or global grad transform
+            joint_forward_backward, _ = _get_gradfn_and_executor(bsym)
+
+            if joint_forward_backward is None:
+                # registered augmented forward impl (where aug fwd and backward are registered separately)
+                aug_fwd_impl = augmented_forward_impls.get(bsym.sym.id)
+                if aug_fwd_impl is not None:
+                    bwd_impl = backward_impls.get(bsym.sym.id)
+
+                    def joint_forward_backward(*args, **kwargs):
+                        aug_fwd_res = aug_fwd_impl(*bsym.args, **bsym.kwargs)
+                        if isinstance(aug_fwd_res, thunder.core.transforms.VJPDual):
+                            res, saved_for_backward = aug_fwd_res
+                        else:
+                            res, saved_for_backward = aug_fwd_res
+                        grad_outs = []
+                        for r in thunder.core.pytree.tree_iter(res):
+                            if isinstance(r, thunder.TensorProxy):
+                                grad_outs.append(prims.get_grad(r))
+                        # to do non-grad outputs of bwd
+                        grad_inps = bwd_impl(*saved_for_backward, *grad_outs)
+                        if isinstance(grad_inps, thunder.Proxy):
+                            grad_inps = [grad_inps]
+
+                        flat_inps = args
+
+                        for i, gi in zip(flat_inps, grad_inps, strict=True):  # strict?
+                            if isinstance(i, thunder.TensorProxy):
+                                prims.put_grad(i, gi)
+                        return res
+
+            if joint_forward_backward is not None:
+                self.new_trace.push_scope([])
+                result = joint_forward_backward(*bsym.args, **bsym.kwargs)
+                self.set_result(result)
+                new_bsyms = self.new_trace.pop_scope()
+
+                forward_part_bsyms = []
+                backward_part_bsyms = []
+                backward_proxy_names = set()
+                for nbsym in new_bsyms:
+                    argnames = {p.name for p in nbsym.flat_proxy_args}
+                    if nbsym.sym in {prims.get_grad, prims.put_grad} or any(
+                        n in backward_proxy_names for n in argnames
+                    ):
+                        backward_proxy_names.update(p.name for p in nbsym.flat_proxy_outs if p.name not in argnames)
+                        backward_part_bsyms.append(nbsym)
+                    else:
+                        forward_part_bsyms.append(nbsym)
+
+                self.add_processed_bsyms(forward_part_bsyms)
+
+                self.collected_bw_part_bsyms[:0] = backward_part_bsyms
+                return
+
+            # No gradient transform found, need to descend
+            check(not bsym.sym.is_prim, lambda: f"Failed to find a gradient transform for bound symbol {bsym=}")
+
+            # TODO: check if this is needed: the old impl checked whether len(bsym.subsymbols) > 0 except for the special case "torch.nn.functional.dropout" with p=0...
+
+            # OUTPUTS to map
+            self.add_unprocessed_bsyms(bsym.subsymbols[:])
+
+    start_time_ns = time.perf_counter_ns()
+
+    processor = AugmentedForwardProcessor(trace)
+    trace, _ = processor()
+    # trace = thunder.core.transform_common.dce(trace)
+
+    end_time_ns = time.perf_counter_ns()
+    elapsed_time_ns = end_time_ns - start_time_ns
+    elapsed_time_millis = elapsed_time_ns // 1000000
+    trace.set_provenance(
+        thunder.core.trace.TraceProvenance(f"Grad transform pass (took {elapsed_time_millis} milliseconds)")
+    )
+    return trace
+
+
+def split_into_forward_and_backward(joint_trace):
+    forward_part_bsyms = []
+    backward_part_bsyms = []
+
+    return_bsym = joint_trace.bound_symbols[-1]
+    assert return_bsym.sym == prims.python_return
+
+    fw_output = return_bsym.args[0]["output"]
+    assert isinstance(fw_output, tuple)
+
+    grad_outs = [None for _ in fw_output]
+
+    output_pos = {o.name: i for i, o in enumerate(fw_output) if isinstance(o, thunder.TensorProxy)}
+    forward_proxy_names = {o.name for o in thunder.core.pytree.tree_iter(fw_output) if isinstance(o, thunder.Proxy)}
+
+    for bsym in reversed(joint_trace.bound_symbols):
+        if bsym.sym == prims.python_return:
+            continue
+
+        if any((o.name in forward_proxy_names) for o in bsym.flat_proxy_outs):
+            forward_part_bsyms.insert(0, bsym.from_bsym())
+            forward_proxy_names.update(a.name for a in bsym.flat_proxy_args)
+            continue
+
+        if bsym.sym == prims.get_grad:
+            grad_outs[output_pos[bsym.args[0].name]] = bsym.output
+            continue
+
+        backward_part_bsyms.insert(0, bsym.from_bsym())
+
+    # collect needed computation
+    saved_for_backward = {}
+    for bsym in backward_part_bsyms:
+        saved_for_backward.update(
+            (a.name, a)
+            for a in bsym.flat_proxy_args
+            if a.name in forward_proxy_names and a.name not in backward_part_bsyms
path: thunder/transforms/autodiff.py
user: login: t-vi
body: Probably does not make a difference, I wanted to not add things already present, but the resulting dict probably is the same (including order).
created_at: 2025-04-15T15:53:05Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,311 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+)
+from thunder.core.utils import check
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                # This is big (and a big messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols in order
+                for bw_bsyms in self.collected_bw_part_bsyms:
+                    # if we don't have and grad_outs, we know that the grad_ins are not needed, so we don't compute them
+                    # (e.g. for `where(a * b > 0, a, b)` we do not need the grad components of a * b propagated
+                    if any(
+                        ((gg.name in grad_proxy_map) or (gg.name in output_proxy_names))
+                        for gg in (bs.args[0] for bs in bw_bsyms)
+                    ):
+                        for nbsym in bw_bsyms:
+                            if nbsym.sym == prims.put_grad:
+                                # put_grad adds something to the gradients
+                                # - if it is an output of the function, we also need to get the
+                                #   "grad_output" (but only once)
+                                # - if we already had a gradient, we need to add the new component
+                                # - structurally (for correctness), there should be no put_grad
+                                #   after a get_grad
+                                p = self.env.get(nbsym.args[0].name, nbsym.args[0])
+                                current_grad = grad_proxy_map.get(p.name)
+                                new_grad = nbsym.args[1]
+                                if current_grad is None and p.name in output_proxy_names:
+                                    self.new_trace.push_scope([])
+                                    current_grad = prims.get_grad(p)
+                                    self.add_processed_bsyms(self.new_trace.pop_scope())
+
+                                if current_grad is not None:
+                                    new_grad = self.add_bsyms_from_function(thunder.torch.add, current_grad, new_grad)
+
+                                grad_proxy_map[p.name] = new_grad
+                                self.write(new_grad, new_grad)
+                            elif nbsym.sym == prims.get_grad:
+                                p = nbsym.args[0]
+                                # TODO: set requires grad here!
+                                name = p.name
+                                current_grad = grad_proxy_map.get(name)
+                                if current_grad is not None:
+                                    # do we also need to map?
+                                    self.write(nbsym.output, current_grad)
+                                    self.swap_map[thunder.core.proxies.variableify(nbsym.output)] = current_grad
+                                    # XXX_replace_output_with_current_grad
+                                elif name in output_proxy_names:
+                                    new_grad = self.add_processed_bsyms([nbsym.from_bsym()])
+                                    grad_proxy_map[name] = new_grad
+                                # elif KNOWN proxy:
+                                #    ### TODO: it can happen that an intermediate does not have put_grad, e.g. first arg to where
+                                else:
+                                    raise RuntimeError(f"grad of non-output and non-intermediate {name} requested")
+                            else:
+                                self.add_processed_bsyms([nbsym.from_bsym()])
+                self.collected_bw_part_bsyms.clear()
+
+                grad_flat_args = []
+                for p in bsym.args[0]["flat_args"]:
+                    # or p = self.read(p) here?
+                    if isinstance(p, thunder.Proxy) and p.requires_grad and p.name in grad_proxy_map:
+                        # is it always OK if we don't have a gradient? (one case: unused input)
+                        # result of put_grad???
+                        grad_flat_args.append(grad_proxy_map[p.name])
+                    else:
+                        grad_flat_args.append(None)
+
+                bsym.args[0]["grad_flat_args"] = grad_flat_args
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+
+            # constant for gradients (no grad required)
+            if is_constant_for_vjp(bsym):
+                if bsym.sym.name == "synchronize":
+                    # This is a *really* terrible hack to cope with non-grad-needing sharded tensors
+                    # as required by LoRA.
+                    from thunder.distributed.prims import all_gather
+
+                    def synchronize_impl(a, group):
+                        return all_gather(a, group, True).wait()
+
+                    self.add_bsyms_from_function(synchronize_impl, *bsym.args, **bsym.kwargs)
+                    return
+
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+
+            # executor or global grad transform
+            joint_forward_backward, _ = _get_gradfn_and_executor(bsym)
+
+            if joint_forward_backward is None:
+                # registered augmented forward impl (where aug fwd and backward are registered separately)
+                aug_fwd_impl = augmented_forward_impls.get(bsym.sym.id)
+                if aug_fwd_impl is not None:
+                    bwd_impl = backward_impls.get(bsym.sym.id)
+
+                    def joint_forward_backward(*args, **kwargs):
+                        aug_fwd_res = aug_fwd_impl(*bsym.args, **bsym.kwargs)
+                        if isinstance(aug_fwd_res, thunder.core.transforms.VJPDual):
+                            res, saved_for_backward = aug_fwd_res
+                        else:
+                            res, saved_for_backward = aug_fwd_res
path: thunder/transforms/autodiff.py
user: login: beverlylytle
body: This looks a little funny.
created_at: 2025-04-16T16:06:38Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,311 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+)
+from thunder.core.utils import check
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                # This is big (and a big messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols in order
+                for bw_bsyms in self.collected_bw_part_bsyms:
+                    # if we don't have and grad_outs, we know that the grad_ins are not needed, so we don't compute them
+                    # (e.g. for `where(a * b > 0, a, b)` we do not need the grad components of a * b propagated
+                    if any(
+                        ((gg.name in grad_proxy_map) or (gg.name in output_proxy_names))
+                        for gg in (bs.args[0] for bs in bw_bsyms)
+                    ):
+                        for nbsym in bw_bsyms:
+                            if nbsym.sym == prims.put_grad:
+                                # put_grad adds something to the gradients
+                                # - if it is an output of the function, we also need to get the
+                                #   "grad_output" (but only once)
+                                # - if we already had a gradient, we need to add the new component
+                                # - structurally (for correctness), there should be no put_grad
+                                #   after a get_grad
+                                p = self.env.get(nbsym.args[0].name, nbsym.args[0])
+                                current_grad = grad_proxy_map.get(p.name)
+                                new_grad = nbsym.args[1]
+                                if current_grad is None and p.name in output_proxy_names:
+                                    self.new_trace.push_scope([])
+                                    current_grad = prims.get_grad(p)
+                                    self.add_processed_bsyms(self.new_trace.pop_scope())
+
+                                if current_grad is not None:
+                                    new_grad = self.add_bsyms_from_function(thunder.torch.add, current_grad, new_grad)
+
+                                grad_proxy_map[p.name] = new_grad
+                                self.write(new_grad, new_grad)
+                            elif nbsym.sym == prims.get_grad:
+                                p = nbsym.args[0]
+                                # TODO: set requires grad here!
+                                name = p.name
+                                current_grad = grad_proxy_map.get(name)
+                                if current_grad is not None:
+                                    # do we also need to map?
+                                    self.write(nbsym.output, current_grad)
+                                    self.swap_map[thunder.core.proxies.variableify(nbsym.output)] = current_grad
+                                    # XXX_replace_output_with_current_grad
+                                elif name in output_proxy_names:
+                                    new_grad = self.add_processed_bsyms([nbsym.from_bsym()])
+                                    grad_proxy_map[name] = new_grad
+                                # elif KNOWN proxy:
+                                #    ### TODO: it can happen that an intermediate does not have put_grad, e.g. first arg to where
+                                else:
+                                    raise RuntimeError(f"grad of non-output and non-intermediate {name} requested")
+                            else:
+                                self.add_processed_bsyms([nbsym.from_bsym()])
+                self.collected_bw_part_bsyms.clear()
+
+                grad_flat_args = []
+                for p in bsym.args[0]["flat_args"]:
+                    # or p = self.read(p) here?
+                    if isinstance(p, thunder.Proxy) and p.requires_grad and p.name in grad_proxy_map:
+                        # is it always OK if we don't have a gradient? (one case: unused input)
+                        # result of put_grad???
+                        grad_flat_args.append(grad_proxy_map[p.name])
+                    else:
+                        grad_flat_args.append(None)
+
+                bsym.args[0]["grad_flat_args"] = grad_flat_args
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+
+            # constant for gradients (no grad required)
+            if is_constant_for_vjp(bsym):
+                if bsym.sym.name == "synchronize":
+                    # This is a *really* terrible hack to cope with non-grad-needing sharded tensors
+                    # as required by LoRA.
+                    from thunder.distributed.prims import all_gather
+
+                    def synchronize_impl(a, group):
+                        return all_gather(a, group, True).wait()
+
+                    self.add_bsyms_from_function(synchronize_impl, *bsym.args, **bsym.kwargs)
+                    return
+
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+
+            # executor or global grad transform
+            joint_forward_backward, _ = _get_gradfn_and_executor(bsym)
+
+            if joint_forward_backward is None:
+                # registered augmented forward impl (where aug fwd and backward are registered separately)
+                aug_fwd_impl = augmented_forward_impls.get(bsym.sym.id)
+                if aug_fwd_impl is not None:
+                    bwd_impl = backward_impls.get(bsym.sym.id)
+
+                    def joint_forward_backward(*args, **kwargs):
+                        aug_fwd_res = aug_fwd_impl(*bsym.args, **bsym.kwargs)
+                        if isinstance(aug_fwd_res, thunder.core.transforms.VJPDual):
+                            res, saved_for_backward = aug_fwd_res
+                        else:
+                            res, saved_for_backward = aug_fwd_res
path: thunder/transforms/autodiff.py
user: login: t-vi
body: I need to figure out what the return convention for aug forward actually is.

Also for backward:

https://github.com/Lightning-AI/lightning-thunder/blob/150d7e0bebda9c2f27bddcf7e21206d0c7d37ee9/thunder/core/transforms.py#L2274-L2280

why is it not `return None, g_weight`

And not to mention that we also have this:

https://github.com/Lightning-AI/lightning-thunder/blob/150d7e0bebda9c2f27bddcf7e21206d0c7d37ee9/thunder/core/transforms.py#L1360-L1377

which looks fishy to me (shouldn't the `put_grad` be on `weight` and `a` is the index?).

created_at: 2025-04-17T14:44:23Z
author_association: COLLABORATOR
diff_hunk: @@ -321,9 +321,8 @@ def __call__(self, *args, **kwargs):
             # vjp transform (applied later).
             def tag_tensorproxy_output_as_detached(proxy):
                 if isinstance(proxy, TensorProxy):
-                    # We need to remove name from trace, otherwise replace will return a proxy with new name.
-                    trace.names.remove(proxy.name)
-                    return proxy.replace(tags=(ProxyTag.DETACHED_AUTOGRAD_GRAPH,))
+                    proxy.tags.add(ProxyTag.DETACHED_AUTOGRAD_GRAPH)
path: thunder/core/symbol.py
user: login: t-vi
body: This is a bit of a funny change, I am seeing the traces for autograd functions having two versions of the proxy, but have not managed to isolate the problem in the old autodiff setting. Apparently it re-executes things more so it does not hit the mismatch in any of the checked traces.

I am relatively confident though, that the trick here is problematic in that it introduces a collision - two distinct proxies with the same name.
created_at: 2025-05-19T19:52:02Z
author_association: COLLABORATOR
diff_hunk: @@ -1944,35 +1944,6 @@ def func(x):
     torch.testing.assert_close(actual_gr, expected_gr)
 
 
-@pytest.mark.parametrize("device", ("cuda", "cpu"))
-def test_backward_recomputation_decomposed_ops(device):
path: thunder/tests/test_grad.py
user: login: t-vi
body: This test is for an old experimental feature `auto_recompute_intermediates` that has not caught on, so I'm removing the test here.
created_at: 2025-05-19T19:52:47Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,311 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+)
+from thunder.core.utils import check
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                # This is big (and a big messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols in order
+                for bw_bsyms in self.collected_bw_part_bsyms:
+                    # if we don't have and grad_outs, we know that the grad_ins are not needed, so we don't compute them
+                    # (e.g. for `where(a * b > 0, a, b)` we do not need the grad components of a * b propagated
+                    if any(
+                        ((gg.name in grad_proxy_map) or (gg.name in output_proxy_names))
+                        for gg in (bs.args[0] for bs in bw_bsyms)
+                    ):
+                        for nbsym in bw_bsyms:
+                            if nbsym.sym == prims.put_grad:
+                                # put_grad adds something to the gradients
+                                # - if it is an output of the function, we also need to get the
+                                #   "grad_output" (but only once)
+                                # - if we already had a gradient, we need to add the new component
+                                # - structurally (for correctness), there should be no put_grad
+                                #   after a get_grad
+                                p = self.env.get(nbsym.args[0].name, nbsym.args[0])
+                                current_grad = grad_proxy_map.get(p.name)
+                                new_grad = nbsym.args[1]
+                                if current_grad is None and p.name in output_proxy_names:
+                                    self.new_trace.push_scope([])
+                                    current_grad = prims.get_grad(p)
+                                    self.add_processed_bsyms(self.new_trace.pop_scope())
+
+                                if current_grad is not None:
+                                    new_grad = self.add_bsyms_from_function(thunder.torch.add, current_grad, new_grad)
+
+                                grad_proxy_map[p.name] = new_grad
+                                self.write(new_grad, new_grad)
+                            elif nbsym.sym == prims.get_grad:
+                                p = nbsym.args[0]
+                                # TODO: set requires grad here!
+                                name = p.name
+                                current_grad = grad_proxy_map.get(name)
+                                if current_grad is not None:
+                                    # do we also need to map?
+                                    self.write(nbsym.output, current_grad)
+                                    self.swap_map[thunder.core.proxies.variableify(nbsym.output)] = current_grad
+                                    # XXX_replace_output_with_current_grad
+                                elif name in output_proxy_names:
+                                    new_grad = self.add_processed_bsyms([nbsym.from_bsym()])
+                                    grad_proxy_map[name] = new_grad
+                                # elif KNOWN proxy:
+                                #    ### TODO: it can happen that an intermediate does not have put_grad, e.g. first arg to where
+                                else:
+                                    raise RuntimeError(f"grad of non-output and non-intermediate {name} requested")
+                            else:
+                                self.add_processed_bsyms([nbsym.from_bsym()])
+                self.collected_bw_part_bsyms.clear()
+
+                grad_flat_args = []
+                for p in bsym.args[0]["flat_args"]:
+                    # or p = self.read(p) here?
+                    if isinstance(p, thunder.Proxy) and p.requires_grad and p.name in grad_proxy_map:
+                        # is it always OK if we don't have a gradient? (one case: unused input)
+                        # result of put_grad???
+                        grad_flat_args.append(grad_proxy_map[p.name])
+                    else:
+                        grad_flat_args.append(None)
+
+                bsym.args[0]["grad_flat_args"] = grad_flat_args
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+
+            # constant for gradients (no grad required)
+            if is_constant_for_vjp(bsym):
+                if bsym.sym.name == "synchronize":
+                    # This is a *really* terrible hack to cope with non-grad-needing sharded tensors
+                    # as required by LoRA.
+                    from thunder.distributed.prims import all_gather
+
+                    def synchronize_impl(a, group):
+                        return all_gather(a, group, True).wait()
+
+                    self.add_bsyms_from_function(synchronize_impl, *bsym.args, **bsym.kwargs)
+                    return
+
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+
+            # executor or global grad transform
+            joint_forward_backward, _ = _get_gradfn_and_executor(bsym)
+
+            if joint_forward_backward is None:
+                # registered augmented forward impl (where aug fwd and backward are registered separately)
+                aug_fwd_impl = augmented_forward_impls.get(bsym.sym.id)
+                if aug_fwd_impl is not None:
+                    bwd_impl = backward_impls.get(bsym.sym.id)
+
+                    def joint_forward_backward(*args, **kwargs):
+                        aug_fwd_res = aug_fwd_impl(*bsym.args, **bsym.kwargs)
+                        if isinstance(aug_fwd_res, thunder.core.transforms.VJPDual):
+                            res, saved_for_backward = aug_fwd_res
+                        else:
+                            res, saved_for_backward = aug_fwd_res
path: thunder/transforms/autodiff.py
user: login: t-vi
body: fixed now.
created_at: 2025-05-19T19:54:19Z
author_association: COLLABORATOR
diff_hunk: @@ -1856,7 +1856,7 @@ def f(x):
 
     with pytest.raises(
         RuntimeError,
-        match="number of outputs of the original forward function must be the same as the number of primal outputs",
+        match="The number of outputs of the gradient transform function must be the same as the nunmber of outputs of the original forward function.",
path: thunder/tests/test_grad.py
user: login: lantiga
body: ```suggestion
        match="The number of outputs of the gradient transform function must be the same as the number of outputs of the original forward function.",
```
created_at: 2025-05-20T11:11:32Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,415 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims, utils
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+    recompute_saved_for_backward,
+)
+from thunder.core.proxies import ProxyTag
+from thunder.core.symbol import BoundSymbol
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+def _should_recompute_bsym_in_backward(bsym):
+    return any((ProxyTag.RECOMPUTE_IN_BACKWARD in o.tags) for o in bsym.flat_proxy_outs)
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                # This is big (and a big messy):
path: thunder/transforms/autodiff.py
user: login: lantiga
body: ```suggestion
                # This is big (and a bit messy):
```
created_at: 2025-05-20T11:15:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,415 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims, utils
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+    recompute_saved_for_backward,
+)
+from thunder.core.proxies import ProxyTag
+from thunder.core.symbol import BoundSymbol
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+def _should_recompute_bsym_in_backward(bsym):
+    return any((ProxyTag.RECOMPUTE_IN_BACKWARD in o.tags) for o in bsym.flat_proxy_outs)
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                # This is big (and a big messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
path: thunder/transforms/autodiff.py
user: login: lantiga
body: ```suggestion
                # The backward computation in self.collected_bw_part_syms is in the order of the
```
created_at: 2025-05-20T11:16:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,477 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims, utils
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+    recompute_saved_for_backward,
+)
+from thunder.core.proxies import ProxyTag
+from thunder.core.symbol import BoundSymbol
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+def _should_recompute_bsym_in_backward(bsym):
+    return any((ProxyTag.RECOMPUTE_IN_BACKWARD in o.tags) for o in bsym.flat_proxy_outs)
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                ########### BEGINNING of return handling (and putting the backward computation in the joint trace)
+                # This is big (and a bit messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
path: thunder/transforms/autodiff.py
user: login: lantiga
body: ```suggestion
                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in torch.autograd.Function)
```
created_at: 2025-05-20T11:21:42Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,477 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims, utils
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+    recompute_saved_for_backward,
+)
+from thunder.core.proxies import ProxyTag
+from thunder.core.symbol import BoundSymbol
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+def _should_recompute_bsym_in_backward(bsym):
+    return any((ProxyTag.RECOMPUTE_IN_BACKWARD in o.tags) for o in bsym.flat_proxy_outs)
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                ########### BEGINNING of return handling (and putting the backward computation in the joint trace)
+                # This is big (and a bit messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols backward in order (from the generation of the list, this
+                # is the backward of the last forward instruction first, then then the backward of the second to last etc.
+                for bw_bsyms in self.collected_bw_part_bsyms:
+                    # if we don't have and grad_outs, we know that the grad_ins are not needed, so we don't compute them
+                    # (e.g. for `where(a * b > 0, a, b)` we do not need the grad components of a * b propagated
+                    if any(
+                        ((gg.name in grad_proxy_map) or (gg.name in output_proxy_names))
+                        for gg in (bs.args[0] for bs in bw_bsyms if bs.sym == prims.get_grad)
+                    ):
+                        for nbsym in bw_bsyms:
+                            if nbsym.sym == prims.put_grad:
+                                # put_grad adds something to the gradients
+                                # - if it is an output of the function, we also need to get the
+                                #   "grad_output" (but only once)
+                                # - if we already had a gradient, we need to add the new component
+                                # - structurally (for correctness), there should be no put_grad
path: thunder/transforms/autodiff.py
user: login: lantiga
body: should we emit an error if that's the case?
created_at: 2025-05-20T11:27:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,477 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims, utils
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+    recompute_saved_for_backward,
+)
+from thunder.core.proxies import ProxyTag
+from thunder.core.symbol import BoundSymbol
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+def _should_recompute_bsym_in_backward(bsym):
+    return any((ProxyTag.RECOMPUTE_IN_BACKWARD in o.tags) for o in bsym.flat_proxy_outs)
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                ########### BEGINNING of return handling (and putting the backward computation in the joint trace)
+                # This is big (and a bit messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols backward in order (from the generation of the list, this
+                # is the backward of the last forward instruction first, then then the backward of the second to last etc.
+                for bw_bsyms in self.collected_bw_part_bsyms:
+                    # if we don't have and grad_outs, we know that the grad_ins are not needed, so we don't compute them
+                    # (e.g. for `where(a * b > 0, a, b)` we do not need the grad components of a * b propagated
+                    if any(
+                        ((gg.name in grad_proxy_map) or (gg.name in output_proxy_names))
+                        for gg in (bs.args[0] for bs in bw_bsyms if bs.sym == prims.get_grad)
+                    ):
+                        for nbsym in bw_bsyms:
+                            if nbsym.sym == prims.put_grad:
+                                # put_grad adds something to the gradients
+                                # - if it is an output of the function, we also need to get the
+                                #   "grad_output" (but only once)
+                                # - if we already had a gradient, we need to add the new component
+                                # - structurally (for correctness), there should be no put_grad
+                                #   after a get_grad
+                                p = self.env.get(nbsym.args[0].name, nbsym.args[0])
+                                current_grad = grad_proxy_map.get(p.name)
+                                new_grad = nbsym.args[1]
+                                if current_grad is None and p.name in output_proxy_names:
+                                    self.new_trace.push_scope([])
path: thunder/transforms/autodiff.py
user: login: lantiga
body: not for now, but we could probably have a helper for the push_scope / something / pop_scope pattern
created_at: 2025-05-20T11:28:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,477 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims, utils
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+    recompute_saved_for_backward,
+)
+from thunder.core.proxies import ProxyTag
+from thunder.core.symbol import BoundSymbol
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+def _should_recompute_bsym_in_backward(bsym):
+    return any((ProxyTag.RECOMPUTE_IN_BACKWARD in o.tags) for o in bsym.flat_proxy_outs)
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                ########### BEGINNING of return handling (and putting the backward computation in the joint trace)
+                # This is big (and a bit messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols backward in order (from the generation of the list, this
+                # is the backward of the last forward instruction first, then then the backward of the second to last etc.
+                for bw_bsyms in self.collected_bw_part_bsyms:
+                    # if we don't have and grad_outs, we know that the grad_ins are not needed, so we don't compute them
+                    # (e.g. for `where(a * b > 0, a, b)` we do not need the grad components of a * b propagated
+                    if any(
+                        ((gg.name in grad_proxy_map) or (gg.name in output_proxy_names))
+                        for gg in (bs.args[0] for bs in bw_bsyms if bs.sym == prims.get_grad)
+                    ):
+                        for nbsym in bw_bsyms:
+                            if nbsym.sym == prims.put_grad:
+                                # put_grad adds something to the gradients
+                                # - if it is an output of the function, we also need to get the
+                                #   "grad_output" (but only once)
+                                # - if we already had a gradient, we need to add the new component
+                                # - structurally (for correctness), there should be no put_grad
+                                #   after a get_grad
+                                p = self.env.get(nbsym.args[0].name, nbsym.args[0])
+                                current_grad = grad_proxy_map.get(p.name)
+                                new_grad = nbsym.args[1]
+                                if current_grad is None and p.name in output_proxy_names:
+                                    self.new_trace.push_scope([])
+                                    current_grad = prims.get_grad(p)
+                                    self.add_processed_bsyms(self.new_trace.pop_scope())
+
+                                if current_grad is not None:
+                                    new_grad = self.add_bsyms_from_function(thunder.torch.add, current_grad, new_grad)
+
+                                grad_proxy_map[p.name] = new_grad
+                                self.write(new_grad, new_grad)
+                            elif nbsym.sym == prims.get_grad:
+                                # get grad reads a gradient, this could be
+                                # - the gradient of an intermediate (possibly combined with an output) form put_grad above
+                                #   note that all put_grads are happen before all get_grads from the dataflow of the forward,
+                                #   in this case we use the gradient from the put_grad(s)
+                                # - an output of the forward function (that has not been combined with an intermediate in pu grd),
+                                #   then the gradient is like the "grad_out" in the parameter list of a autograd.Function.backward
+                                #   in this case we keep a get_grad in the trace.
+                                # - an intermedite that does not have a gradient computed. In this case the gradient is 0 in the shape
+                                #   of the get_grad argument. This happens e.g. if we have something like "a > 0" for a float a then
+                                #   a has grad 0.
+                                #   TODO: We could (and should) optimize this (see below).
+                                p = nbsym.args[0]
+
+                                # TODO: set proxy requires grad here!
+                                name = p.name
+                                current_grad = grad_proxy_map.get(name)
+                                if current_grad is not None:
+                                    # do we also need to map?
+                                    self.write(nbsym.output, current_grad)
+                                    # replace_output_with_current_grad
+                                    self.swap_map[thunder.core.proxies.variableify(nbsym.output)] = current_grad
+                                elif name in output_proxy_names:
+                                    # output here???
+                                    new_bsym = nbsym.from_bsym()
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                else:
+                                    # TODO: mark this? if all inputs to the backward formula are unused, we would not want to compute it.
+                                    new_bsym = thunder.torch.zeros.bind(
+                                        *p.shape, device=p.device, dtype=p.dtype, output=nbsym.output
+                                    )
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                    self.write(nbsym.output, new_bsym.output)
+                            else:
+                                # all other symbols are just computation
+                                self.add_processed_bsyms([nbsym.from_bsym()])
+                self.collected_bw_part_bsyms.clear()
+
+                # we collect the gradients fo the flat args in the output dictionary's 'grad_flat_args'
+                # this means that the joint trace as them as outputs, which is good.
+                grad_flat_args = []
+                for p in bsym.args[0]["flat_args"]:
+                    # or p = self.read(p) here?
+                    if isinstance(p, thunder.TensorProxy) and p.requires_grad and p.name in grad_proxy_map:
+                        # is it always OK if we don't have a gradient? (one case: unused input)
+                        # result of put_grad???
+                        grad_flat_args.append(grad_proxy_map[p.name])
+                    else:
+                        grad_flat_args.append(None)
+
+                bsym.args[0]["grad_flat_args"] = grad_flat_args
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+                ########### END of return handling (and putting the backward computation in the joint trace)
+
+            # we now handle various bound symbols by collecting augmented forward and backward symbols
+
+            # 1. constant for gradients (no grad required)
+            #    excemting synchronize here terrible hack to cope with non-grad-needing sharded tensors
path: thunder/transforms/autodiff.py
user: login: lantiga
body: ```suggestion
            #    executing synchronize here terrible hack to cope with non-grad-needing sharded tensors
```
probably we could extend the explanation a bit here
created_at: 2025-05-20T11:30:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,477 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims, utils
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+    recompute_saved_for_backward,
+)
+from thunder.core.proxies import ProxyTag
+from thunder.core.symbol import BoundSymbol
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+def _should_recompute_bsym_in_backward(bsym):
+    return any((ProxyTag.RECOMPUTE_IN_BACKWARD in o.tags) for o in bsym.flat_proxy_outs)
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                ########### BEGINNING of return handling (and putting the backward computation in the joint trace)
+                # This is big (and a bit messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols backward in order (from the generation of the list, this
+                # is the backward of the last forward instruction first, then then the backward of the second to last etc.
+                for bw_bsyms in self.collected_bw_part_bsyms:
+                    # if we don't have and grad_outs, we know that the grad_ins are not needed, so we don't compute them
+                    # (e.g. for `where(a * b > 0, a, b)` we do not need the grad components of a * b propagated
+                    if any(
+                        ((gg.name in grad_proxy_map) or (gg.name in output_proxy_names))
+                        for gg in (bs.args[0] for bs in bw_bsyms if bs.sym == prims.get_grad)
+                    ):
+                        for nbsym in bw_bsyms:
+                            if nbsym.sym == prims.put_grad:
+                                # put_grad adds something to the gradients
+                                # - if it is an output of the function, we also need to get the
+                                #   "grad_output" (but only once)
+                                # - if we already had a gradient, we need to add the new component
+                                # - structurally (for correctness), there should be no put_grad
+                                #   after a get_grad
+                                p = self.env.get(nbsym.args[0].name, nbsym.args[0])
+                                current_grad = grad_proxy_map.get(p.name)
+                                new_grad = nbsym.args[1]
+                                if current_grad is None and p.name in output_proxy_names:
+                                    self.new_trace.push_scope([])
+                                    current_grad = prims.get_grad(p)
+                                    self.add_processed_bsyms(self.new_trace.pop_scope())
+
+                                if current_grad is not None:
+                                    new_grad = self.add_bsyms_from_function(thunder.torch.add, current_grad, new_grad)
+
+                                grad_proxy_map[p.name] = new_grad
+                                self.write(new_grad, new_grad)
+                            elif nbsym.sym == prims.get_grad:
+                                # get grad reads a gradient, this could be
+                                # - the gradient of an intermediate (possibly combined with an output) form put_grad above
+                                #   note that all put_grads are happen before all get_grads from the dataflow of the forward,
+                                #   in this case we use the gradient from the put_grad(s)
+                                # - an output of the forward function (that has not been combined with an intermediate in pu grd),
+                                #   then the gradient is like the "grad_out" in the parameter list of a autograd.Function.backward
+                                #   in this case we keep a get_grad in the trace.
+                                # - an intermedite that does not have a gradient computed. In this case the gradient is 0 in the shape
+                                #   of the get_grad argument. This happens e.g. if we have something like "a > 0" for a float a then
+                                #   a has grad 0.
+                                #   TODO: We could (and should) optimize this (see below).
+                                p = nbsym.args[0]
+
+                                # TODO: set proxy requires grad here!
+                                name = p.name
+                                current_grad = grad_proxy_map.get(name)
+                                if current_grad is not None:
+                                    # do we also need to map?
+                                    self.write(nbsym.output, current_grad)
+                                    # replace_output_with_current_grad
+                                    self.swap_map[thunder.core.proxies.variableify(nbsym.output)] = current_grad
+                                elif name in output_proxy_names:
+                                    # output here???
+                                    new_bsym = nbsym.from_bsym()
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                else:
+                                    # TODO: mark this? if all inputs to the backward formula are unused, we would not want to compute it.
+                                    new_bsym = thunder.torch.zeros.bind(
+                                        *p.shape, device=p.device, dtype=p.dtype, output=nbsym.output
+                                    )
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                    self.write(nbsym.output, new_bsym.output)
+                            else:
+                                # all other symbols are just computation
+                                self.add_processed_bsyms([nbsym.from_bsym()])
+                self.collected_bw_part_bsyms.clear()
+
+                # we collect the gradients fo the flat args in the output dictionary's 'grad_flat_args'
+                # this means that the joint trace as them as outputs, which is good.
+                grad_flat_args = []
+                for p in bsym.args[0]["flat_args"]:
+                    # or p = self.read(p) here?
+                    if isinstance(p, thunder.TensorProxy) and p.requires_grad and p.name in grad_proxy_map:
+                        # is it always OK if we don't have a gradient? (one case: unused input)
+                        # result of put_grad???
+                        grad_flat_args.append(grad_proxy_map[p.name])
+                    else:
+                        grad_flat_args.append(None)
+
+                bsym.args[0]["grad_flat_args"] = grad_flat_args
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+                ########### END of return handling (and putting the backward computation in the joint trace)
+
+            # we now handle various bound symbols by collecting augmented forward and backward symbols
+
+            # 1. constant for gradients (no grad required)
+            #    excemting synchronize here terrible hack to cope with non-grad-needing sharded tensors
+            #    as required by LoRA. We should have a symbol tag "always compute grad" instead.
+            if is_constant_for_vjp(bsym) and not bsym.sym.name == "synchronize":
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+
+            # 2. Special case the thunder.torch.checkpoint higher oder function
path: thunder/transforms/autodiff.py
user: login: lantiga
body: ```suggestion
            # 2. Special case the thunder.torch.checkpoint higher order function
```
created_at: 2025-05-20T11:31:30Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,477 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims, utils
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+    recompute_saved_for_backward,
+)
+from thunder.core.proxies import ProxyTag
+from thunder.core.symbol import BoundSymbol
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+def _should_recompute_bsym_in_backward(bsym):
+    return any((ProxyTag.RECOMPUTE_IN_BACKWARD in o.tags) for o in bsym.flat_proxy_outs)
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                ########### BEGINNING of return handling (and putting the backward computation in the joint trace)
+                # This is big (and a bit messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols backward in order (from the generation of the list, this
+                # is the backward of the last forward instruction first, then then the backward of the second to last etc.
+                for bw_bsyms in self.collected_bw_part_bsyms:
+                    # if we don't have and grad_outs, we know that the grad_ins are not needed, so we don't compute them
+                    # (e.g. for `where(a * b > 0, a, b)` we do not need the grad components of a * b propagated
+                    if any(
+                        ((gg.name in grad_proxy_map) or (gg.name in output_proxy_names))
+                        for gg in (bs.args[0] for bs in bw_bsyms if bs.sym == prims.get_grad)
+                    ):
+                        for nbsym in bw_bsyms:
+                            if nbsym.sym == prims.put_grad:
+                                # put_grad adds something to the gradients
+                                # - if it is an output of the function, we also need to get the
+                                #   "grad_output" (but only once)
+                                # - if we already had a gradient, we need to add the new component
+                                # - structurally (for correctness), there should be no put_grad
+                                #   after a get_grad
+                                p = self.env.get(nbsym.args[0].name, nbsym.args[0])
+                                current_grad = grad_proxy_map.get(p.name)
+                                new_grad = nbsym.args[1]
+                                if current_grad is None and p.name in output_proxy_names:
+                                    self.new_trace.push_scope([])
+                                    current_grad = prims.get_grad(p)
+                                    self.add_processed_bsyms(self.new_trace.pop_scope())
+
+                                if current_grad is not None:
+                                    new_grad = self.add_bsyms_from_function(thunder.torch.add, current_grad, new_grad)
+
+                                grad_proxy_map[p.name] = new_grad
+                                self.write(new_grad, new_grad)
+                            elif nbsym.sym == prims.get_grad:
+                                # get grad reads a gradient, this could be
+                                # - the gradient of an intermediate (possibly combined with an output) form put_grad above
+                                #   note that all put_grads are happen before all get_grads from the dataflow of the forward,
+                                #   in this case we use the gradient from the put_grad(s)
+                                # - an output of the forward function (that has not been combined with an intermediate in pu grd),
+                                #   then the gradient is like the "grad_out" in the parameter list of a autograd.Function.backward
+                                #   in this case we keep a get_grad in the trace.
+                                # - an intermedite that does not have a gradient computed. In this case the gradient is 0 in the shape
+                                #   of the get_grad argument. This happens e.g. if we have something like "a > 0" for a float a then
+                                #   a has grad 0.
+                                #   TODO: We could (and should) optimize this (see below).
+                                p = nbsym.args[0]
+
+                                # TODO: set proxy requires grad here!
+                                name = p.name
+                                current_grad = grad_proxy_map.get(name)
+                                if current_grad is not None:
+                                    # do we also need to map?
+                                    self.write(nbsym.output, current_grad)
+                                    # replace_output_with_current_grad
+                                    self.swap_map[thunder.core.proxies.variableify(nbsym.output)] = current_grad
+                                elif name in output_proxy_names:
+                                    # output here???
+                                    new_bsym = nbsym.from_bsym()
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                else:
+                                    # TODO: mark this? if all inputs to the backward formula are unused, we would not want to compute it.
+                                    new_bsym = thunder.torch.zeros.bind(
+                                        *p.shape, device=p.device, dtype=p.dtype, output=nbsym.output
+                                    )
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                    self.write(nbsym.output, new_bsym.output)
+                            else:
+                                # all other symbols are just computation
+                                self.add_processed_bsyms([nbsym.from_bsym()])
+                self.collected_bw_part_bsyms.clear()
+
+                # we collect the gradients fo the flat args in the output dictionary's 'grad_flat_args'
+                # this means that the joint trace as them as outputs, which is good.
+                grad_flat_args = []
+                for p in bsym.args[0]["flat_args"]:
+                    # or p = self.read(p) here?
+                    if isinstance(p, thunder.TensorProxy) and p.requires_grad and p.name in grad_proxy_map:
+                        # is it always OK if we don't have a gradient? (one case: unused input)
+                        # result of put_grad???
+                        grad_flat_args.append(grad_proxy_map[p.name])
+                    else:
+                        grad_flat_args.append(None)
+
+                bsym.args[0]["grad_flat_args"] = grad_flat_args
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+                ########### END of return handling (and putting the backward computation in the joint trace)
+
+            # we now handle various bound symbols by collecting augmented forward and backward symbols
+
+            # 1. constant for gradients (no grad required)
+            #    excemting synchronize here terrible hack to cope with non-grad-needing sharded tensors
+            #    as required by LoRA. We should have a symbol tag "always compute grad" instead.
+            if is_constant_for_vjp(bsym) and not bsym.sym.name == "synchronize":
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+
+            # 2. Special case the thunder.torch.checkpoint higher oder function
+            if bsym.sym == thunder.torch.checkpoint:
+                # Tag all intermediate outputs as to be recomputed.
+                function_arg_names = {a.name for a in bsym.flat_proxy_args}
+
+                for subsym in bsym.subsymbols:
+                    for o in subsym.flat_proxy_outs:
+                        if o.name not in function_arg_names:
+                            o.tags.add(ProxyTag.RECOMPUTE_IN_BACKWARD)
+
+                # decompose
+                self.add_unprocessed_bsyms(bsym.subsymbols[:])
+                return
+
+            # 3a. see if we have a grad_transform (e.g. from OperatorExecutor.register_grad_transform)
+
+            # executor or global grad transform
+            joint_forward_backward, _ = _get_gradfn_and_executor(bsym)
+
+            if joint_forward_backward is None:
+                # 3b. if we don't have a grad_transform for a bsym, maybe we have old style augmented forward and backward
+                # registered. In this case, we make the joint_forward_backward which works like a grad_gransform.
+                # registered augmented forward impl (where aug fwd and backward are registered separately)
+                aug_fwd_impl = augmented_forward_impls.get(bsym.sym.id)
+                if aug_fwd_impl is not None:
+                    bwd_impl = backward_impls.get(bsym.sym.id)
+
+                    # this is our ad hoc combined forward and backward function ("grad_transform")
+                    def joint_forward_backward(*args, **kwargs):
+                        arg_proxy_names = {a.name for a in bsym.flat_proxy_args}
+
+                        # run the augmented forward res
+                        aug_fwd_res = aug_fwd_impl(*bsym.args, **bsym.kwargs)
+
+                        # aug_fwd_res could be either VJPDual or just a tuple, in any case, we can decompose it.
+                        res, saved_for_backward = aug_fwd_res
+
+                        # we need to shallow copy inputs that are returned for "get_grad" and "put_grad" to properly work
+                        # (this shallow copy is the equivalent because we of an "edge" in the PyTorch autograd graph)
+                        def shallow_copy_if_input(p):
+                            if isinstance(p, thunder.TensorProxy) and p.name in arg_proxy_names:
+                                return thunder.core.prims.shallow_copy(p)
+                            return p
+
+                        res = tree_map(shallow_copy_if_input, res)
+
+                        # now we need the backward. it starts by getting the grad_outs
+                        grad_outs = []
+                        for r in thunder.core.pytree.tree_iter(res):
+                            if isinstance(r, thunder.TensorProxy):
+                                grad_outs.append(prims.get_grad(r))
+
+                        # The backward computes the grad_inps of the bsym from the grad_outs
+                        # TODO: non-grad outputs of bwd?
+                        grad_inps = bwd_impl(*saved_for_backward, *grad_outs)
+                        if isinstance(grad_inps, thunder.Proxy):
+                            grad_inps = [grad_inps]
+
+                        # match the grad_inps to the inputs of the boudnd symbol and put the grads
+
+                        flat_inps = args
+                        # for autograd_function_apply, skip the function args
+                        # TODO: fix the returned gradients to include two None?.
+                        if bsym.sym == thunder.torch.autograd_function_apply:
+                            flat_inps = args[2:]
+
+                        # there may be non-gradient requiring additional args (todo: maybe only support this for non-tensor ones?)
+                        num_flat_tensor_inps = sum(isinstance(i, thunder.TensorProxy) for i in flat_inps)
+                        utils.check(
+                            num_flat_tensor_inps <= len(grad_inps),
+                            lambda: f"Backward for {bsym.sym.id} returned {len(grad_inps)} value(s), but expected {num_flat_tensor_inps}",
+                        )
+
+                        assert len(grad_inps) <= len(flat_inps)
+                        for i, gi in zip(flat_inps, grad_inps):
+                            # for integer proxies etc. we expect gi to be None
+                            if isinstance(i, thunder.TensorProxy) and gi is not None:
+                                prims.put_grad(i, gi)
+                        return res
+
+            # 3c if we have a joint forward backward (either grad_transform or the ad hoc constructed one above), we can comptue the gradient
+            if joint_forward_backward is not None:
+                self.new_trace.push_scope([])
+                result = joint_forward_backward(*bsym.args, **bsym.kwargs)
+
+                # Check that the number of outputs of the original forward function is the
+                # same as the number of primal outputs of the augmented forward trace
+                utils.check(
+                    len(utils.sequencify(bsym.output)) == len(utils.sequencify(result)),
+                    lambda: f"While generating forward and backward functions for {bsym.sym.name}, encountered an error.\n"
+                    "The number of outputs of the gradient transform function must be the same as the nunmber of outputs of the original forward function.\n"
+                    f"Number of outputs of the original forward function: {len(utils.sequencify(bsym.output))}\n"
+                    f"Number of primal outputs of the gradient tansform / augmented forward: {len(utils.sequencify(result))}\n"
+                    "Please check the forward function and the gradient transfrm function / augmented forward to ensure that they have the same number of outputs.",
path: thunder/transforms/autodiff.py
user: login: lantiga
body: ```suggestion
                    "Please check the forward function and the gradient transform function / augmented forward to ensure that they have the same number of outputs.",
```
created_at: 2025-05-20T11:35:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,477 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims, utils
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+    recompute_saved_for_backward,
+)
+from thunder.core.proxies import ProxyTag
+from thunder.core.symbol import BoundSymbol
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+def _should_recompute_bsym_in_backward(bsym):
+    return any((ProxyTag.RECOMPUTE_IN_BACKWARD in o.tags) for o in bsym.flat_proxy_outs)
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                ########### BEGINNING of return handling (and putting the backward computation in the joint trace)
+                # This is big (and a bit messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols backward in order (from the generation of the list, this
+                # is the backward of the last forward instruction first, then then the backward of the second to last etc.
+                for bw_bsyms in self.collected_bw_part_bsyms:
+                    # if we don't have and grad_outs, we know that the grad_ins are not needed, so we don't compute them
+                    # (e.g. for `where(a * b > 0, a, b)` we do not need the grad components of a * b propagated
+                    if any(
+                        ((gg.name in grad_proxy_map) or (gg.name in output_proxy_names))
+                        for gg in (bs.args[0] for bs in bw_bsyms if bs.sym == prims.get_grad)
+                    ):
+                        for nbsym in bw_bsyms:
+                            if nbsym.sym == prims.put_grad:
+                                # put_grad adds something to the gradients
+                                # - if it is an output of the function, we also need to get the
+                                #   "grad_output" (but only once)
+                                # - if we already had a gradient, we need to add the new component
+                                # - structurally (for correctness), there should be no put_grad
+                                #   after a get_grad
+                                p = self.env.get(nbsym.args[0].name, nbsym.args[0])
+                                current_grad = grad_proxy_map.get(p.name)
+                                new_grad = nbsym.args[1]
+                                if current_grad is None and p.name in output_proxy_names:
+                                    self.new_trace.push_scope([])
+                                    current_grad = prims.get_grad(p)
+                                    self.add_processed_bsyms(self.new_trace.pop_scope())
+
+                                if current_grad is not None:
+                                    new_grad = self.add_bsyms_from_function(thunder.torch.add, current_grad, new_grad)
+
+                                grad_proxy_map[p.name] = new_grad
+                                self.write(new_grad, new_grad)
+                            elif nbsym.sym == prims.get_grad:
+                                # get grad reads a gradient, this could be
+                                # - the gradient of an intermediate (possibly combined with an output) form put_grad above
+                                #   note that all put_grads are happen before all get_grads from the dataflow of the forward,
+                                #   in this case we use the gradient from the put_grad(s)
+                                # - an output of the forward function (that has not been combined with an intermediate in pu grd),
+                                #   then the gradient is like the "grad_out" in the parameter list of a autograd.Function.backward
+                                #   in this case we keep a get_grad in the trace.
+                                # - an intermedite that does not have a gradient computed. In this case the gradient is 0 in the shape
+                                #   of the get_grad argument. This happens e.g. if we have something like "a > 0" for a float a then
+                                #   a has grad 0.
+                                #   TODO: We could (and should) optimize this (see below).
+                                p = nbsym.args[0]
+
+                                # TODO: set proxy requires grad here!
+                                name = p.name
+                                current_grad = grad_proxy_map.get(name)
+                                if current_grad is not None:
+                                    # do we also need to map?
+                                    self.write(nbsym.output, current_grad)
+                                    # replace_output_with_current_grad
+                                    self.swap_map[thunder.core.proxies.variableify(nbsym.output)] = current_grad
+                                elif name in output_proxy_names:
+                                    # output here???
+                                    new_bsym = nbsym.from_bsym()
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                else:
+                                    # TODO: mark this? if all inputs to the backward formula are unused, we would not want to compute it.
+                                    new_bsym = thunder.torch.zeros.bind(
+                                        *p.shape, device=p.device, dtype=p.dtype, output=nbsym.output
+                                    )
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                    self.write(nbsym.output, new_bsym.output)
+                            else:
+                                # all other symbols are just computation
+                                self.add_processed_bsyms([nbsym.from_bsym()])
+                self.collected_bw_part_bsyms.clear()
+
+                # we collect the gradients fo the flat args in the output dictionary's 'grad_flat_args'
+                # this means that the joint trace as them as outputs, which is good.
+                grad_flat_args = []
+                for p in bsym.args[0]["flat_args"]:
+                    # or p = self.read(p) here?
+                    if isinstance(p, thunder.TensorProxy) and p.requires_grad and p.name in grad_proxy_map:
+                        # is it always OK if we don't have a gradient? (one case: unused input)
+                        # result of put_grad???
+                        grad_flat_args.append(grad_proxy_map[p.name])
+                    else:
+                        grad_flat_args.append(None)
+
+                bsym.args[0]["grad_flat_args"] = grad_flat_args
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+                ########### END of return handling (and putting the backward computation in the joint trace)
+
+            # we now handle various bound symbols by collecting augmented forward and backward symbols
+
+            # 1. constant for gradients (no grad required)
+            #    excemting synchronize here terrible hack to cope with non-grad-needing sharded tensors
+            #    as required by LoRA. We should have a symbol tag "always compute grad" instead.
+            if is_constant_for_vjp(bsym) and not bsym.sym.name == "synchronize":
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+
+            # 2. Special case the thunder.torch.checkpoint higher oder function
+            if bsym.sym == thunder.torch.checkpoint:
+                # Tag all intermediate outputs as to be recomputed.
+                function_arg_names = {a.name for a in bsym.flat_proxy_args}
+
+                for subsym in bsym.subsymbols:
+                    for o in subsym.flat_proxy_outs:
+                        if o.name not in function_arg_names:
+                            o.tags.add(ProxyTag.RECOMPUTE_IN_BACKWARD)
+
+                # decompose
+                self.add_unprocessed_bsyms(bsym.subsymbols[:])
+                return
+
+            # 3a. see if we have a grad_transform (e.g. from OperatorExecutor.register_grad_transform)
+
+            # executor or global grad transform
+            joint_forward_backward, _ = _get_gradfn_and_executor(bsym)
+
+            if joint_forward_backward is None:
+                # 3b. if we don't have a grad_transform for a bsym, maybe we have old style augmented forward and backward
+                # registered. In this case, we make the joint_forward_backward which works like a grad_gransform.
+                # registered augmented forward impl (where aug fwd and backward are registered separately)
+                aug_fwd_impl = augmented_forward_impls.get(bsym.sym.id)
+                if aug_fwd_impl is not None:
+                    bwd_impl = backward_impls.get(bsym.sym.id)
+
+                    # this is our ad hoc combined forward and backward function ("grad_transform")
+                    def joint_forward_backward(*args, **kwargs):
+                        arg_proxy_names = {a.name for a in bsym.flat_proxy_args}
+
+                        # run the augmented forward res
+                        aug_fwd_res = aug_fwd_impl(*bsym.args, **bsym.kwargs)
+
+                        # aug_fwd_res could be either VJPDual or just a tuple, in any case, we can decompose it.
+                        res, saved_for_backward = aug_fwd_res
+
+                        # we need to shallow copy inputs that are returned for "get_grad" and "put_grad" to properly work
+                        # (this shallow copy is the equivalent because we of an "edge" in the PyTorch autograd graph)
+                        def shallow_copy_if_input(p):
+                            if isinstance(p, thunder.TensorProxy) and p.name in arg_proxy_names:
+                                return thunder.core.prims.shallow_copy(p)
+                            return p
+
+                        res = tree_map(shallow_copy_if_input, res)
+
+                        # now we need the backward. it starts by getting the grad_outs
+                        grad_outs = []
+                        for r in thunder.core.pytree.tree_iter(res):
+                            if isinstance(r, thunder.TensorProxy):
+                                grad_outs.append(prims.get_grad(r))
+
+                        # The backward computes the grad_inps of the bsym from the grad_outs
+                        # TODO: non-grad outputs of bwd?
+                        grad_inps = bwd_impl(*saved_for_backward, *grad_outs)
+                        if isinstance(grad_inps, thunder.Proxy):
+                            grad_inps = [grad_inps]
+
+                        # match the grad_inps to the inputs of the boudnd symbol and put the grads
+
+                        flat_inps = args
+                        # for autograd_function_apply, skip the function args
+                        # TODO: fix the returned gradients to include two None?.
+                        if bsym.sym == thunder.torch.autograd_function_apply:
+                            flat_inps = args[2:]
+
+                        # there may be non-gradient requiring additional args (todo: maybe only support this for non-tensor ones?)
+                        num_flat_tensor_inps = sum(isinstance(i, thunder.TensorProxy) for i in flat_inps)
+                        utils.check(
+                            num_flat_tensor_inps <= len(grad_inps),
+                            lambda: f"Backward for {bsym.sym.id} returned {len(grad_inps)} value(s), but expected {num_flat_tensor_inps}",
+                        )
+
+                        assert len(grad_inps) <= len(flat_inps)
+                        for i, gi in zip(flat_inps, grad_inps):
+                            # for integer proxies etc. we expect gi to be None
+                            if isinstance(i, thunder.TensorProxy) and gi is not None:
+                                prims.put_grad(i, gi)
+                        return res
+
+            # 3c if we have a joint forward backward (either grad_transform or the ad hoc constructed one above), we can comptue the gradient
+            if joint_forward_backward is not None:
+                self.new_trace.push_scope([])
+                result = joint_forward_backward(*bsym.args, **bsym.kwargs)
+
+                # Check that the number of outputs of the original forward function is the
+                # same as the number of primal outputs of the augmented forward trace
+                utils.check(
+                    len(utils.sequencify(bsym.output)) == len(utils.sequencify(result)),
+                    lambda: f"While generating forward and backward functions for {bsym.sym.name}, encountered an error.\n"
+                    "The number of outputs of the gradient transform function must be the same as the nunmber of outputs of the original forward function.\n"
+                    f"Number of outputs of the original forward function: {len(utils.sequencify(bsym.output))}\n"
+                    f"Number of primal outputs of the gradient tansform / augmented forward: {len(utils.sequencify(result))}\n"
+                    "Please check the forward function and the gradient transfrm function / augmented forward to ensure that they have the same number of outputs.",
+                )
+                self.set_result(result)
+                new_bsyms = self.new_trace.pop_scope()
+
+                # simple splitting: only compute in forward what is needed for the output
+                forward_part_proxy_names = {
+                    o.name for o in thunder.core.pytree.tree_iter(result) if isinstance(o, thunder.Proxy)
+                }
+                forward_part_bsyms = []
+                backward_part_bsyms = []
+                for nbsym in reversed(new_bsyms):
+                    # argnames = {p.name for p in nbsym.flat_proxy_args}
+                    if nbsym.sym in {prims.get_grad, prims.put_grad}:
+                        backward_part_bsyms.insert(0, nbsym)
+                        continue
+                    assert nbsym.sym != prims.unpack_trivial  # can unpack trivial happen here?
+                    if any((o.name in forward_part_proxy_names) for o in nbsym.flat_proxy_outs):
+                        forward_part_bsyms.insert(0, nbsym)
+                        bsym_arg_names = {a.name for a in nbsym.flat_proxy_args}
+                        bsym_out_names = {o.name for o in nbsym.flat_proxy_outs if o.name not in bsym_arg_names}
+                        forward_part_proxy_names.update(bsym_arg_names)
+                        forward_part_proxy_names.update(bsym_out_names)
+                        continue
+                    backward_part_bsyms.insert(0, nbsym)
+                self.add_processed_bsyms(forward_part_bsyms)
+
+                self.collected_bw_part_bsyms.insert(0, backward_part_bsyms)
+                return
+
+            # 4. No gradient transform found, need to decompose the symbol
+
+            # error if this is a primitive
+            utils.check(not bsym.sym.is_prim, lambda: f"Failed to find a gradient transform for bound symbol {bsym=}")
+
+            # TODO: check if this is needed: the old impl checked whether len(bsym.subsymbols) > 0 except for the special case "torch.nn.functional.dropout" with p=0...
+            # add the decomposition (= the subsymbols) to the front of the symbols to be processed
+            self.add_unprocessed_bsyms(bsym.subsymbols[:])
+
+            # end of 4 and end of the bsym processing loop.
+
+    start_time_ns = time.perf_counter_ns()
+
+    # run the trace through the processor
+    trace, _ = AugmentedForwardProcessor(trace)()
+    # run through DCE in case some of the gradients of intermediates are not needed.
+    trace = thunder.core.transform_common.dce(trace)
+
+    end_time_ns = time.perf_counter_ns()
+    elapsed_time_ns = end_time_ns - start_time_ns
+    elapsed_time_millis = elapsed_time_ns // 1000000
+    trace.set_provenance(
+        thunder.core.trace.TraceProvenance(f"Grad transform pass (took {elapsed_time_millis} milliseconds)")
+    )
+    return trace
+
+
+def split_into_forward_and_backward(joint_trace):
+    """split a joint trace for forward and backward into separate ones, including recomputation (aka activation checkpointing)"""
+
+    # the joint trace will have the forward computation at the beginning and then the backward computation
+    # from how it is constructed.
+    # we split the trace:
+    # - forward symbols go into forward_part_bsyms
+    # - all symbols not in the forward go into backward_part_bsyms
+    # - for recomputation (aka activation checkpointing), we want to insert symbols going into the forward also into the
+    #   backward. but we want to do so "just in time". To this end, we gather the symbols in a dict and later
path: thunder/transforms/autodiff.py
user: login: lantiga
body: ```suggestion
    #   backward, but we want to do so "just in time". To this end, we gather the symbols in a dict and later
```
created_at: 2025-05-20T11:38:26Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,477 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims, utils
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+    recompute_saved_for_backward,
+)
+from thunder.core.proxies import ProxyTag
+from thunder.core.symbol import BoundSymbol
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+def _should_recompute_bsym_in_backward(bsym):
+    return any((ProxyTag.RECOMPUTE_IN_BACKWARD in o.tags) for o in bsym.flat_proxy_outs)
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                ########### BEGINNING of return handling (and putting the backward computation in the joint trace)
+                # This is big (and a bit messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols backward in order (from the generation of the list, this
+                # is the backward of the last forward instruction first, then then the backward of the second to last etc.
+                for bw_bsyms in self.collected_bw_part_bsyms:
+                    # if we don't have and grad_outs, we know that the grad_ins are not needed, so we don't compute them
+                    # (e.g. for `where(a * b > 0, a, b)` we do not need the grad components of a * b propagated
+                    if any(
+                        ((gg.name in grad_proxy_map) or (gg.name in output_proxy_names))
+                        for gg in (bs.args[0] for bs in bw_bsyms if bs.sym == prims.get_grad)
+                    ):
+                        for nbsym in bw_bsyms:
+                            if nbsym.sym == prims.put_grad:
+                                # put_grad adds something to the gradients
+                                # - if it is an output of the function, we also need to get the
+                                #   "grad_output" (but only once)
+                                # - if we already had a gradient, we need to add the new component
+                                # - structurally (for correctness), there should be no put_grad
+                                #   after a get_grad
+                                p = self.env.get(nbsym.args[0].name, nbsym.args[0])
+                                current_grad = grad_proxy_map.get(p.name)
+                                new_grad = nbsym.args[1]
+                                if current_grad is None and p.name in output_proxy_names:
+                                    self.new_trace.push_scope([])
+                                    current_grad = prims.get_grad(p)
+                                    self.add_processed_bsyms(self.new_trace.pop_scope())
+
+                                if current_grad is not None:
+                                    new_grad = self.add_bsyms_from_function(thunder.torch.add, current_grad, new_grad)
+
+                                grad_proxy_map[p.name] = new_grad
+                                self.write(new_grad, new_grad)
+                            elif nbsym.sym == prims.get_grad:
+                                # get grad reads a gradient, this could be
+                                # - the gradient of an intermediate (possibly combined with an output) form put_grad above
+                                #   note that all put_grads are happen before all get_grads from the dataflow of the forward,
+                                #   in this case we use the gradient from the put_grad(s)
+                                # - an output of the forward function (that has not been combined with an intermediate in pu grd),
+                                #   then the gradient is like the "grad_out" in the parameter list of a autograd.Function.backward
+                                #   in this case we keep a get_grad in the trace.
+                                # - an intermedite that does not have a gradient computed. In this case the gradient is 0 in the shape
+                                #   of the get_grad argument. This happens e.g. if we have something like "a > 0" for a float a then
+                                #   a has grad 0.
+                                #   TODO: We could (and should) optimize this (see below).
+                                p = nbsym.args[0]
+
+                                # TODO: set proxy requires grad here!
+                                name = p.name
+                                current_grad = grad_proxy_map.get(name)
+                                if current_grad is not None:
+                                    # do we also need to map?
+                                    self.write(nbsym.output, current_grad)
+                                    # replace_output_with_current_grad
+                                    self.swap_map[thunder.core.proxies.variableify(nbsym.output)] = current_grad
+                                elif name in output_proxy_names:
+                                    # output here???
+                                    new_bsym = nbsym.from_bsym()
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                else:
+                                    # TODO: mark this? if all inputs to the backward formula are unused, we would not want to compute it.
+                                    new_bsym = thunder.torch.zeros.bind(
+                                        *p.shape, device=p.device, dtype=p.dtype, output=nbsym.output
+                                    )
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                    self.write(nbsym.output, new_bsym.output)
+                            else:
+                                # all other symbols are just computation
+                                self.add_processed_bsyms([nbsym.from_bsym()])
+                self.collected_bw_part_bsyms.clear()
+
+                # we collect the gradients fo the flat args in the output dictionary's 'grad_flat_args'
+                # this means that the joint trace as them as outputs, which is good.
+                grad_flat_args = []
+                for p in bsym.args[0]["flat_args"]:
+                    # or p = self.read(p) here?
+                    if isinstance(p, thunder.TensorProxy) and p.requires_grad and p.name in grad_proxy_map:
+                        # is it always OK if we don't have a gradient? (one case: unused input)
+                        # result of put_grad???
+                        grad_flat_args.append(grad_proxy_map[p.name])
+                    else:
+                        grad_flat_args.append(None)
+
+                bsym.args[0]["grad_flat_args"] = grad_flat_args
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+                ########### END of return handling (and putting the backward computation in the joint trace)
+
+            # we now handle various bound symbols by collecting augmented forward and backward symbols
+
+            # 1. constant for gradients (no grad required)
+            #    excemting synchronize here terrible hack to cope with non-grad-needing sharded tensors
+            #    as required by LoRA. We should have a symbol tag "always compute grad" instead.
+            if is_constant_for_vjp(bsym) and not bsym.sym.name == "synchronize":
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+
+            # 2. Special case the thunder.torch.checkpoint higher oder function
+            if bsym.sym == thunder.torch.checkpoint:
+                # Tag all intermediate outputs as to be recomputed.
+                function_arg_names = {a.name for a in bsym.flat_proxy_args}
+
+                for subsym in bsym.subsymbols:
+                    for o in subsym.flat_proxy_outs:
+                        if o.name not in function_arg_names:
+                            o.tags.add(ProxyTag.RECOMPUTE_IN_BACKWARD)
+
+                # decompose
+                self.add_unprocessed_bsyms(bsym.subsymbols[:])
+                return
+
+            # 3a. see if we have a grad_transform (e.g. from OperatorExecutor.register_grad_transform)
+
+            # executor or global grad transform
+            joint_forward_backward, _ = _get_gradfn_and_executor(bsym)
+
+            if joint_forward_backward is None:
+                # 3b. if we don't have a grad_transform for a bsym, maybe we have old style augmented forward and backward
+                # registered. In this case, we make the joint_forward_backward which works like a grad_gransform.
+                # registered augmented forward impl (where aug fwd and backward are registered separately)
+                aug_fwd_impl = augmented_forward_impls.get(bsym.sym.id)
+                if aug_fwd_impl is not None:
+                    bwd_impl = backward_impls.get(bsym.sym.id)
+
+                    # this is our ad hoc combined forward and backward function ("grad_transform")
+                    def joint_forward_backward(*args, **kwargs):
+                        arg_proxy_names = {a.name for a in bsym.flat_proxy_args}
+
+                        # run the augmented forward res
+                        aug_fwd_res = aug_fwd_impl(*bsym.args, **bsym.kwargs)
+
+                        # aug_fwd_res could be either VJPDual or just a tuple, in any case, we can decompose it.
+                        res, saved_for_backward = aug_fwd_res
+
+                        # we need to shallow copy inputs that are returned for "get_grad" and "put_grad" to properly work
+                        # (this shallow copy is the equivalent because we of an "edge" in the PyTorch autograd graph)
+                        def shallow_copy_if_input(p):
+                            if isinstance(p, thunder.TensorProxy) and p.name in arg_proxy_names:
+                                return thunder.core.prims.shallow_copy(p)
+                            return p
+
+                        res = tree_map(shallow_copy_if_input, res)
+
+                        # now we need the backward. it starts by getting the grad_outs
+                        grad_outs = []
+                        for r in thunder.core.pytree.tree_iter(res):
+                            if isinstance(r, thunder.TensorProxy):
+                                grad_outs.append(prims.get_grad(r))
+
+                        # The backward computes the grad_inps of the bsym from the grad_outs
+                        # TODO: non-grad outputs of bwd?
+                        grad_inps = bwd_impl(*saved_for_backward, *grad_outs)
+                        if isinstance(grad_inps, thunder.Proxy):
+                            grad_inps = [grad_inps]
+
+                        # match the grad_inps to the inputs of the boudnd symbol and put the grads
+
+                        flat_inps = args
+                        # for autograd_function_apply, skip the function args
+                        # TODO: fix the returned gradients to include two None?.
+                        if bsym.sym == thunder.torch.autograd_function_apply:
+                            flat_inps = args[2:]
+
+                        # there may be non-gradient requiring additional args (todo: maybe only support this for non-tensor ones?)
+                        num_flat_tensor_inps = sum(isinstance(i, thunder.TensorProxy) for i in flat_inps)
+                        utils.check(
+                            num_flat_tensor_inps <= len(grad_inps),
+                            lambda: f"Backward for {bsym.sym.id} returned {len(grad_inps)} value(s), but expected {num_flat_tensor_inps}",
+                        )
+
+                        assert len(grad_inps) <= len(flat_inps)
+                        for i, gi in zip(flat_inps, grad_inps):
+                            # for integer proxies etc. we expect gi to be None
+                            if isinstance(i, thunder.TensorProxy) and gi is not None:
+                                prims.put_grad(i, gi)
+                        return res
+
+            # 3c if we have a joint forward backward (either grad_transform or the ad hoc constructed one above), we can comptue the gradient
+            if joint_forward_backward is not None:
+                self.new_trace.push_scope([])
+                result = joint_forward_backward(*bsym.args, **bsym.kwargs)
+
+                # Check that the number of outputs of the original forward function is the
+                # same as the number of primal outputs of the augmented forward trace
+                utils.check(
+                    len(utils.sequencify(bsym.output)) == len(utils.sequencify(result)),
+                    lambda: f"While generating forward and backward functions for {bsym.sym.name}, encountered an error.\n"
+                    "The number of outputs of the gradient transform function must be the same as the nunmber of outputs of the original forward function.\n"
+                    f"Number of outputs of the original forward function: {len(utils.sequencify(bsym.output))}\n"
+                    f"Number of primal outputs of the gradient tansform / augmented forward: {len(utils.sequencify(result))}\n"
+                    "Please check the forward function and the gradient transfrm function / augmented forward to ensure that they have the same number of outputs.",
+                )
+                self.set_result(result)
+                new_bsyms = self.new_trace.pop_scope()
+
+                # simple splitting: only compute in forward what is needed for the output
+                forward_part_proxy_names = {
+                    o.name for o in thunder.core.pytree.tree_iter(result) if isinstance(o, thunder.Proxy)
+                }
+                forward_part_bsyms = []
+                backward_part_bsyms = []
+                for nbsym in reversed(new_bsyms):
+                    # argnames = {p.name for p in nbsym.flat_proxy_args}
+                    if nbsym.sym in {prims.get_grad, prims.put_grad}:
+                        backward_part_bsyms.insert(0, nbsym)
+                        continue
+                    assert nbsym.sym != prims.unpack_trivial  # can unpack trivial happen here?
+                    if any((o.name in forward_part_proxy_names) for o in nbsym.flat_proxy_outs):
+                        forward_part_bsyms.insert(0, nbsym)
+                        bsym_arg_names = {a.name for a in nbsym.flat_proxy_args}
+                        bsym_out_names = {o.name for o in nbsym.flat_proxy_outs if o.name not in bsym_arg_names}
+                        forward_part_proxy_names.update(bsym_arg_names)
+                        forward_part_proxy_names.update(bsym_out_names)
+                        continue
+                    backward_part_bsyms.insert(0, nbsym)
+                self.add_processed_bsyms(forward_part_bsyms)
+
+                self.collected_bw_part_bsyms.insert(0, backward_part_bsyms)
+                return
+
+            # 4. No gradient transform found, need to decompose the symbol
+
+            # error if this is a primitive
+            utils.check(not bsym.sym.is_prim, lambda: f"Failed to find a gradient transform for bound symbol {bsym=}")
+
+            # TODO: check if this is needed: the old impl checked whether len(bsym.subsymbols) > 0 except for the special case "torch.nn.functional.dropout" with p=0...
+            # add the decomposition (= the subsymbols) to the front of the symbols to be processed
+            self.add_unprocessed_bsyms(bsym.subsymbols[:])
+
+            # end of 4 and end of the bsym processing loop.
+
+    start_time_ns = time.perf_counter_ns()
+
+    # run the trace through the processor
+    trace, _ = AugmentedForwardProcessor(trace)()
+    # run through DCE in case some of the gradients of intermediates are not needed.
+    trace = thunder.core.transform_common.dce(trace)
+
+    end_time_ns = time.perf_counter_ns()
+    elapsed_time_ns = end_time_ns - start_time_ns
+    elapsed_time_millis = elapsed_time_ns // 1000000
+    trace.set_provenance(
+        thunder.core.trace.TraceProvenance(f"Grad transform pass (took {elapsed_time_millis} milliseconds)")
+    )
+    return trace
+
+
+def split_into_forward_and_backward(joint_trace):
+    """split a joint trace for forward and backward into separate ones, including recomputation (aka activation checkpointing)"""
+
+    # the joint trace will have the forward computation at the beginning and then the backward computation
+    # from how it is constructed.
+    # we split the trace:
+    # - forward symbols go into forward_part_bsyms
+    # - all symbols not in the forward go into backward_part_bsyms
+    # - for recomputation (aka activation checkpointing), we want to insert symbols going into the forward also into the
+    #   backward. but we want to do so "just in time". To this end, we gather the symbols in a dict and later
+    #   insert it when their respective outputs are needed. This is in backward_part_bsyms_recomputated
path: thunder/transforms/autodiff.py
user: login: lantiga
body: ```suggestion
    #   insert it when their respective outputs are needed. This is in backward_part_bsyms_recomputed
```
created_at: 2025-05-20T11:39:08Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,477 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims, utils
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+    recompute_saved_for_backward,
+)
+from thunder.core.proxies import ProxyTag
+from thunder.core.symbol import BoundSymbol
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+def _should_recompute_bsym_in_backward(bsym):
+    return any((ProxyTag.RECOMPUTE_IN_BACKWARD in o.tags) for o in bsym.flat_proxy_outs)
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
path: thunder/transforms/autodiff.py
user: login: beverlylytle
body: ```suggestion
    # - if there is a augmented forward registered, use that and the backward registered,
```
created_at: 2025-05-20T12:17:41Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,477 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims, utils
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+    recompute_saved_for_backward,
+)
+from thunder.core.proxies import ProxyTag
+from thunder.core.symbol import BoundSymbol
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+def _should_recompute_bsym_in_backward(bsym):
+    return any((ProxyTag.RECOMPUTE_IN_BACKWARD in o.tags) for o in bsym.flat_proxy_outs)
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                ########### BEGINNING of return handling (and putting the backward computation in the joint trace)
+                # This is big (and a bit messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols backward in order (from the generation of the list, this
+                # is the backward of the last forward instruction first, then then the backward of the second to last etc.
+                for bw_bsyms in self.collected_bw_part_bsyms:
+                    # if we don't have and grad_outs, we know that the grad_ins are not needed, so we don't compute them
+                    # (e.g. for `where(a * b > 0, a, b)` we do not need the grad components of a * b propagated
+                    if any(
+                        ((gg.name in grad_proxy_map) or (gg.name in output_proxy_names))
+                        for gg in (bs.args[0] for bs in bw_bsyms if bs.sym == prims.get_grad)
+                    ):
+                        for nbsym in bw_bsyms:
+                            if nbsym.sym == prims.put_grad:
+                                # put_grad adds something to the gradients
+                                # - if it is an output of the function, we also need to get the
+                                #   "grad_output" (but only once)
+                                # - if we already had a gradient, we need to add the new component
+                                # - structurally (for correctness), there should be no put_grad
+                                #   after a get_grad
+                                p = self.env.get(nbsym.args[0].name, nbsym.args[0])
+                                current_grad = grad_proxy_map.get(p.name)
+                                new_grad = nbsym.args[1]
+                                if current_grad is None and p.name in output_proxy_names:
+                                    self.new_trace.push_scope([])
+                                    current_grad = prims.get_grad(p)
+                                    self.add_processed_bsyms(self.new_trace.pop_scope())
+
+                                if current_grad is not None:
+                                    new_grad = self.add_bsyms_from_function(thunder.torch.add, current_grad, new_grad)
+
+                                grad_proxy_map[p.name] = new_grad
+                                self.write(new_grad, new_grad)
+                            elif nbsym.sym == prims.get_grad:
+                                # get grad reads a gradient, this could be
+                                # - the gradient of an intermediate (possibly combined with an output) form put_grad above
+                                #   note that all put_grads are happen before all get_grads from the dataflow of the forward,
+                                #   in this case we use the gradient from the put_grad(s)
+                                # - an output of the forward function (that has not been combined with an intermediate in pu grd),
path: thunder/transforms/autodiff.py
user: login: beverlylytle
body: "put grad"?
created_at: 2025-05-20T12:19:55Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,477 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims, utils
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+    recompute_saved_for_backward,
+)
+from thunder.core.proxies import ProxyTag
+from thunder.core.symbol import BoundSymbol
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+def _should_recompute_bsym_in_backward(bsym):
+    return any((ProxyTag.RECOMPUTE_IN_BACKWARD in o.tags) for o in bsym.flat_proxy_outs)
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                ########### BEGINNING of return handling (and putting the backward computation in the joint trace)
+                # This is big (and a bit messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols backward in order (from the generation of the list, this
+                # is the backward of the last forward instruction first, then then the backward of the second to last etc.
+                for bw_bsyms in self.collected_bw_part_bsyms:
+                    # if we don't have and grad_outs, we know that the grad_ins are not needed, so we don't compute them
+                    # (e.g. for `where(a * b > 0, a, b)` we do not need the grad components of a * b propagated
+                    if any(
+                        ((gg.name in grad_proxy_map) or (gg.name in output_proxy_names))
+                        for gg in (bs.args[0] for bs in bw_bsyms if bs.sym == prims.get_grad)
+                    ):
+                        for nbsym in bw_bsyms:
+                            if nbsym.sym == prims.put_grad:
+                                # put_grad adds something to the gradients
+                                # - if it is an output of the function, we also need to get the
+                                #   "grad_output" (but only once)
+                                # - if we already had a gradient, we need to add the new component
+                                # - structurally (for correctness), there should be no put_grad
+                                #   after a get_grad
+                                p = self.env.get(nbsym.args[0].name, nbsym.args[0])
+                                current_grad = grad_proxy_map.get(p.name)
+                                new_grad = nbsym.args[1]
+                                if current_grad is None and p.name in output_proxy_names:
+                                    self.new_trace.push_scope([])
+                                    current_grad = prims.get_grad(p)
+                                    self.add_processed_bsyms(self.new_trace.pop_scope())
+
+                                if current_grad is not None:
+                                    new_grad = self.add_bsyms_from_function(thunder.torch.add, current_grad, new_grad)
+
+                                grad_proxy_map[p.name] = new_grad
+                                self.write(new_grad, new_grad)
+                            elif nbsym.sym == prims.get_grad:
+                                # get grad reads a gradient, this could be
+                                # - the gradient of an intermediate (possibly combined with an output) form put_grad above
+                                #   note that all put_grads are happen before all get_grads from the dataflow of the forward,
+                                #   in this case we use the gradient from the put_grad(s)
+                                # - an output of the forward function (that has not been combined with an intermediate in pu grd),
+                                #   then the gradient is like the "grad_out" in the parameter list of a autograd.Function.backward
+                                #   in this case we keep a get_grad in the trace.
+                                # - an intermedite that does not have a gradient computed. In this case the gradient is 0 in the shape
path: thunder/transforms/autodiff.py
user: login: beverlylytle
body: ```suggestion
                                # - an intermediate that does not have a gradient computed. In this case the gradient is 0 in the shape
```
created_at: 2025-05-20T12:20:48Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,477 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims, utils
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+    recompute_saved_for_backward,
+)
+from thunder.core.proxies import ProxyTag
+from thunder.core.symbol import BoundSymbol
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+def _should_recompute_bsym_in_backward(bsym):
+    return any((ProxyTag.RECOMPUTE_IN_BACKWARD in o.tags) for o in bsym.flat_proxy_outs)
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                ########### BEGINNING of return handling (and putting the backward computation in the joint trace)
+                # This is big (and a bit messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols backward in order (from the generation of the list, this
+                # is the backward of the last forward instruction first, then then the backward of the second to last etc.
+                for bw_bsyms in self.collected_bw_part_bsyms:
+                    # if we don't have and grad_outs, we know that the grad_ins are not needed, so we don't compute them
+                    # (e.g. for `where(a * b > 0, a, b)` we do not need the grad components of a * b propagated
+                    if any(
+                        ((gg.name in grad_proxy_map) or (gg.name in output_proxy_names))
+                        for gg in (bs.args[0] for bs in bw_bsyms if bs.sym == prims.get_grad)
+                    ):
+                        for nbsym in bw_bsyms:
+                            if nbsym.sym == prims.put_grad:
+                                # put_grad adds something to the gradients
+                                # - if it is an output of the function, we also need to get the
+                                #   "grad_output" (but only once)
+                                # - if we already had a gradient, we need to add the new component
+                                # - structurally (for correctness), there should be no put_grad
+                                #   after a get_grad
+                                p = self.env.get(nbsym.args[0].name, nbsym.args[0])
+                                current_grad = grad_proxy_map.get(p.name)
+                                new_grad = nbsym.args[1]
+                                if current_grad is None and p.name in output_proxy_names:
+                                    self.new_trace.push_scope([])
+                                    current_grad = prims.get_grad(p)
+                                    self.add_processed_bsyms(self.new_trace.pop_scope())
+
+                                if current_grad is not None:
+                                    new_grad = self.add_bsyms_from_function(thunder.torch.add, current_grad, new_grad)
+
+                                grad_proxy_map[p.name] = new_grad
+                                self.write(new_grad, new_grad)
+                            elif nbsym.sym == prims.get_grad:
+                                # get grad reads a gradient, this could be
+                                # - the gradient of an intermediate (possibly combined with an output) form put_grad above
+                                #   note that all put_grads are happen before all get_grads from the dataflow of the forward,
+                                #   in this case we use the gradient from the put_grad(s)
+                                # - an output of the forward function (that has not been combined with an intermediate in pu grd),
+                                #   then the gradient is like the "grad_out" in the parameter list of a autograd.Function.backward
+                                #   in this case we keep a get_grad in the trace.
+                                # - an intermedite that does not have a gradient computed. In this case the gradient is 0 in the shape
+                                #   of the get_grad argument. This happens e.g. if we have something like "a > 0" for a float a then
+                                #   a has grad 0.
+                                #   TODO: We could (and should) optimize this (see below).
+                                p = nbsym.args[0]
+
+                                # TODO: set proxy requires grad here!
+                                name = p.name
+                                current_grad = grad_proxy_map.get(name)
+                                if current_grad is not None:
+                                    # do we also need to map?
+                                    self.write(nbsym.output, current_grad)
+                                    # replace_output_with_current_grad
+                                    self.swap_map[thunder.core.proxies.variableify(nbsym.output)] = current_grad
+                                elif name in output_proxy_names:
+                                    # output here???
+                                    new_bsym = nbsym.from_bsym()
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                else:
+                                    # TODO: mark this? if all inputs to the backward formula are unused, we would not want to compute it.
+                                    new_bsym = thunder.torch.zeros.bind(
+                                        *p.shape, device=p.device, dtype=p.dtype, output=nbsym.output
+                                    )
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                    self.write(nbsym.output, new_bsym.output)
+                            else:
+                                # all other symbols are just computation
+                                self.add_processed_bsyms([nbsym.from_bsym()])
+                self.collected_bw_part_bsyms.clear()
+
+                # we collect the gradients fo the flat args in the output dictionary's 'grad_flat_args'
+                # this means that the joint trace as them as outputs, which is good.
path: thunder/transforms/autodiff.py
user: login: beverlylytle
body: ```suggestion
                # this means that the joint trace has them as outputs, which is good.
```
created_at: 2025-05-20T12:21:23Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,477 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims, utils
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+    recompute_saved_for_backward,
+)
+from thunder.core.proxies import ProxyTag
+from thunder.core.symbol import BoundSymbol
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+def _should_recompute_bsym_in_backward(bsym):
+    return any((ProxyTag.RECOMPUTE_IN_BACKWARD in o.tags) for o in bsym.flat_proxy_outs)
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                ########### BEGINNING of return handling (and putting the backward computation in the joint trace)
+                # This is big (and a bit messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols backward in order (from the generation of the list, this
+                # is the backward of the last forward instruction first, then then the backward of the second to last etc.
+                for bw_bsyms in self.collected_bw_part_bsyms:
+                    # if we don't have and grad_outs, we know that the grad_ins are not needed, so we don't compute them
+                    # (e.g. for `where(a * b > 0, a, b)` we do not need the grad components of a * b propagated
+                    if any(
+                        ((gg.name in grad_proxy_map) or (gg.name in output_proxy_names))
+                        for gg in (bs.args[0] for bs in bw_bsyms if bs.sym == prims.get_grad)
+                    ):
+                        for nbsym in bw_bsyms:
+                            if nbsym.sym == prims.put_grad:
+                                # put_grad adds something to the gradients
+                                # - if it is an output of the function, we also need to get the
+                                #   "grad_output" (but only once)
+                                # - if we already had a gradient, we need to add the new component
+                                # - structurally (for correctness), there should be no put_grad
+                                #   after a get_grad
+                                p = self.env.get(nbsym.args[0].name, nbsym.args[0])
+                                current_grad = grad_proxy_map.get(p.name)
+                                new_grad = nbsym.args[1]
+                                if current_grad is None and p.name in output_proxy_names:
+                                    self.new_trace.push_scope([])
+                                    current_grad = prims.get_grad(p)
+                                    self.add_processed_bsyms(self.new_trace.pop_scope())
+
+                                if current_grad is not None:
+                                    new_grad = self.add_bsyms_from_function(thunder.torch.add, current_grad, new_grad)
+
+                                grad_proxy_map[p.name] = new_grad
+                                self.write(new_grad, new_grad)
+                            elif nbsym.sym == prims.get_grad:
+                                # get grad reads a gradient, this could be
+                                # - the gradient of an intermediate (possibly combined with an output) form put_grad above
+                                #   note that all put_grads are happen before all get_grads from the dataflow of the forward,
+                                #   in this case we use the gradient from the put_grad(s)
+                                # - an output of the forward function (that has not been combined with an intermediate in pu grd),
+                                #   then the gradient is like the "grad_out" in the parameter list of a autograd.Function.backward
+                                #   in this case we keep a get_grad in the trace.
+                                # - an intermedite that does not have a gradient computed. In this case the gradient is 0 in the shape
+                                #   of the get_grad argument. This happens e.g. if we have something like "a > 0" for a float a then
+                                #   a has grad 0.
+                                #   TODO: We could (and should) optimize this (see below).
+                                p = nbsym.args[0]
+
+                                # TODO: set proxy requires grad here!
+                                name = p.name
+                                current_grad = grad_proxy_map.get(name)
+                                if current_grad is not None:
+                                    # do we also need to map?
+                                    self.write(nbsym.output, current_grad)
+                                    # replace_output_with_current_grad
+                                    self.swap_map[thunder.core.proxies.variableify(nbsym.output)] = current_grad
+                                elif name in output_proxy_names:
+                                    # output here???
+                                    new_bsym = nbsym.from_bsym()
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                else:
+                                    # TODO: mark this? if all inputs to the backward formula are unused, we would not want to compute it.
+                                    new_bsym = thunder.torch.zeros.bind(
+                                        *p.shape, device=p.device, dtype=p.dtype, output=nbsym.output
+                                    )
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                    self.write(nbsym.output, new_bsym.output)
+                            else:
+                                # all other symbols are just computation
+                                self.add_processed_bsyms([nbsym.from_bsym()])
+                self.collected_bw_part_bsyms.clear()
+
+                # we collect the gradients fo the flat args in the output dictionary's 'grad_flat_args'
+                # this means that the joint trace as them as outputs, which is good.
+                grad_flat_args = []
+                for p in bsym.args[0]["flat_args"]:
+                    # or p = self.read(p) here?
+                    if isinstance(p, thunder.TensorProxy) and p.requires_grad and p.name in grad_proxy_map:
+                        # is it always OK if we don't have a gradient? (one case: unused input)
+                        # result of put_grad???
+                        grad_flat_args.append(grad_proxy_map[p.name])
+                    else:
+                        grad_flat_args.append(None)
+
+                bsym.args[0]["grad_flat_args"] = grad_flat_args
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+                ########### END of return handling (and putting the backward computation in the joint trace)
+
+            # we now handle various bound symbols by collecting augmented forward and backward symbols
+
+            # 1. constant for gradients (no grad required)
+            #    excemting synchronize here terrible hack to cope with non-grad-needing sharded tensors
+            #    as required by LoRA. We should have a symbol tag "always compute grad" instead.
+            if is_constant_for_vjp(bsym) and not bsym.sym.name == "synchronize":
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+
+            # 2. Special case the thunder.torch.checkpoint higher oder function
+            if bsym.sym == thunder.torch.checkpoint:
+                # Tag all intermediate outputs as to be recomputed.
+                function_arg_names = {a.name for a in bsym.flat_proxy_args}
+
+                for subsym in bsym.subsymbols:
+                    for o in subsym.flat_proxy_outs:
+                        if o.name not in function_arg_names:
+                            o.tags.add(ProxyTag.RECOMPUTE_IN_BACKWARD)
+
+                # decompose
+                self.add_unprocessed_bsyms(bsym.subsymbols[:])
+                return
+
+            # 3a. see if we have a grad_transform (e.g. from OperatorExecutor.register_grad_transform)
+
+            # executor or global grad transform
+            joint_forward_backward, _ = _get_gradfn_and_executor(bsym)
+
+            if joint_forward_backward is None:
+                # 3b. if we don't have a grad_transform for a bsym, maybe we have old style augmented forward and backward
+                # registered. In this case, we make the joint_forward_backward which works like a grad_gransform.
+                # registered augmented forward impl (where aug fwd and backward are registered separately)
+                aug_fwd_impl = augmented_forward_impls.get(bsym.sym.id)
+                if aug_fwd_impl is not None:
+                    bwd_impl = backward_impls.get(bsym.sym.id)
+
+                    # this is our ad hoc combined forward and backward function ("grad_transform")
+                    def joint_forward_backward(*args, **kwargs):
+                        arg_proxy_names = {a.name for a in bsym.flat_proxy_args}
+
+                        # run the augmented forward res
+                        aug_fwd_res = aug_fwd_impl(*bsym.args, **bsym.kwargs)
+
+                        # aug_fwd_res could be either VJPDual or just a tuple, in any case, we can decompose it.
+                        res, saved_for_backward = aug_fwd_res
+
+                        # we need to shallow copy inputs that are returned for "get_grad" and "put_grad" to properly work
+                        # (this shallow copy is the equivalent because we of an "edge" in the PyTorch autograd graph)
+                        def shallow_copy_if_input(p):
+                            if isinstance(p, thunder.TensorProxy) and p.name in arg_proxy_names:
+                                return thunder.core.prims.shallow_copy(p)
+                            return p
+
+                        res = tree_map(shallow_copy_if_input, res)
+
+                        # now we need the backward. it starts by getting the grad_outs
+                        grad_outs = []
+                        for r in thunder.core.pytree.tree_iter(res):
+                            if isinstance(r, thunder.TensorProxy):
+                                grad_outs.append(prims.get_grad(r))
+
+                        # The backward computes the grad_inps of the bsym from the grad_outs
+                        # TODO: non-grad outputs of bwd?
+                        grad_inps = bwd_impl(*saved_for_backward, *grad_outs)
+                        if isinstance(grad_inps, thunder.Proxy):
+                            grad_inps = [grad_inps]
+
+                        # match the grad_inps to the inputs of the boudnd symbol and put the grads
+
+                        flat_inps = args
+                        # for autograd_function_apply, skip the function args
+                        # TODO: fix the returned gradients to include two None?.
+                        if bsym.sym == thunder.torch.autograd_function_apply:
+                            flat_inps = args[2:]
+
+                        # there may be non-gradient requiring additional args (todo: maybe only support this for non-tensor ones?)
+                        num_flat_tensor_inps = sum(isinstance(i, thunder.TensorProxy) for i in flat_inps)
+                        utils.check(
+                            num_flat_tensor_inps <= len(grad_inps),
+                            lambda: f"Backward for {bsym.sym.id} returned {len(grad_inps)} value(s), but expected {num_flat_tensor_inps}",
+                        )
+
+                        assert len(grad_inps) <= len(flat_inps)
+                        for i, gi in zip(flat_inps, grad_inps):
+                            # for integer proxies etc. we expect gi to be None
+                            if isinstance(i, thunder.TensorProxy) and gi is not None:
+                                prims.put_grad(i, gi)
+                        return res
+
+            # 3c if we have a joint forward backward (either grad_transform or the ad hoc constructed one above), we can comptue the gradient
+            if joint_forward_backward is not None:
+                self.new_trace.push_scope([])
+                result = joint_forward_backward(*bsym.args, **bsym.kwargs)
+
+                # Check that the number of outputs of the original forward function is the
+                # same as the number of primal outputs of the augmented forward trace
+                utils.check(
+                    len(utils.sequencify(bsym.output)) == len(utils.sequencify(result)),
+                    lambda: f"While generating forward and backward functions for {bsym.sym.name}, encountered an error.\n"
+                    "The number of outputs of the gradient transform function must be the same as the nunmber of outputs of the original forward function.\n"
path: thunder/transforms/autodiff.py
user: login: beverlylytle
body: ```suggestion
                    "The number of outputs of the gradient transform function must be the same as the number of outputs of the original forward function.\n"
```
created_at: 2025-05-20T12:23:19Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,477 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims, utils
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+    recompute_saved_for_backward,
+)
+from thunder.core.proxies import ProxyTag
+from thunder.core.symbol import BoundSymbol
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+def _should_recompute_bsym_in_backward(bsym):
+    return any((ProxyTag.RECOMPUTE_IN_BACKWARD in o.tags) for o in bsym.flat_proxy_outs)
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                ########### BEGINNING of return handling (and putting the backward computation in the joint trace)
+                # This is big (and a bit messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols backward in order (from the generation of the list, this
+                # is the backward of the last forward instruction first, then then the backward of the second to last etc.
+                for bw_bsyms in self.collected_bw_part_bsyms:
+                    # if we don't have and grad_outs, we know that the grad_ins are not needed, so we don't compute them
+                    # (e.g. for `where(a * b > 0, a, b)` we do not need the grad components of a * b propagated
+                    if any(
+                        ((gg.name in grad_proxy_map) or (gg.name in output_proxy_names))
+                        for gg in (bs.args[0] for bs in bw_bsyms if bs.sym == prims.get_grad)
+                    ):
+                        for nbsym in bw_bsyms:
+                            if nbsym.sym == prims.put_grad:
+                                # put_grad adds something to the gradients
+                                # - if it is an output of the function, we also need to get the
+                                #   "grad_output" (but only once)
+                                # - if we already had a gradient, we need to add the new component
+                                # - structurally (for correctness), there should be no put_grad
+                                #   after a get_grad
+                                p = self.env.get(nbsym.args[0].name, nbsym.args[0])
+                                current_grad = grad_proxy_map.get(p.name)
+                                new_grad = nbsym.args[1]
+                                if current_grad is None and p.name in output_proxy_names:
+                                    self.new_trace.push_scope([])
+                                    current_grad = prims.get_grad(p)
+                                    self.add_processed_bsyms(self.new_trace.pop_scope())
+
+                                if current_grad is not None:
+                                    new_grad = self.add_bsyms_from_function(thunder.torch.add, current_grad, new_grad)
+
+                                grad_proxy_map[p.name] = new_grad
+                                self.write(new_grad, new_grad)
+                            elif nbsym.sym == prims.get_grad:
+                                # get grad reads a gradient, this could be
+                                # - the gradient of an intermediate (possibly combined with an output) form put_grad above
+                                #   note that all put_grads are happen before all get_grads from the dataflow of the forward,
+                                #   in this case we use the gradient from the put_grad(s)
+                                # - an output of the forward function (that has not been combined with an intermediate in pu grd),
+                                #   then the gradient is like the "grad_out" in the parameter list of a autograd.Function.backward
+                                #   in this case we keep a get_grad in the trace.
+                                # - an intermedite that does not have a gradient computed. In this case the gradient is 0 in the shape
+                                #   of the get_grad argument. This happens e.g. if we have something like "a > 0" for a float a then
+                                #   a has grad 0.
+                                #   TODO: We could (and should) optimize this (see below).
+                                p = nbsym.args[0]
+
+                                # TODO: set proxy requires grad here!
+                                name = p.name
+                                current_grad = grad_proxy_map.get(name)
+                                if current_grad is not None:
+                                    # do we also need to map?
+                                    self.write(nbsym.output, current_grad)
+                                    # replace_output_with_current_grad
+                                    self.swap_map[thunder.core.proxies.variableify(nbsym.output)] = current_grad
+                                elif name in output_proxy_names:
+                                    # output here???
+                                    new_bsym = nbsym.from_bsym()
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                else:
+                                    # TODO: mark this? if all inputs to the backward formula are unused, we would not want to compute it.
+                                    new_bsym = thunder.torch.zeros.bind(
+                                        *p.shape, device=p.device, dtype=p.dtype, output=nbsym.output
+                                    )
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                    self.write(nbsym.output, new_bsym.output)
+                            else:
+                                # all other symbols are just computation
+                                self.add_processed_bsyms([nbsym.from_bsym()])
+                self.collected_bw_part_bsyms.clear()
+
+                # we collect the gradients fo the flat args in the output dictionary's 'grad_flat_args'
+                # this means that the joint trace as them as outputs, which is good.
+                grad_flat_args = []
+                for p in bsym.args[0]["flat_args"]:
+                    # or p = self.read(p) here?
+                    if isinstance(p, thunder.TensorProxy) and p.requires_grad and p.name in grad_proxy_map:
+                        # is it always OK if we don't have a gradient? (one case: unused input)
+                        # result of put_grad???
+                        grad_flat_args.append(grad_proxy_map[p.name])
+                    else:
+                        grad_flat_args.append(None)
+
+                bsym.args[0]["grad_flat_args"] = grad_flat_args
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+                ########### END of return handling (and putting the backward computation in the joint trace)
+
+            # we now handle various bound symbols by collecting augmented forward and backward symbols
+
+            # 1. constant for gradients (no grad required)
+            #    excemting synchronize here terrible hack to cope with non-grad-needing sharded tensors
+            #    as required by LoRA. We should have a symbol tag "always compute grad" instead.
+            if is_constant_for_vjp(bsym) and not bsym.sym.name == "synchronize":
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+
+            # 2. Special case the thunder.torch.checkpoint higher oder function
+            if bsym.sym == thunder.torch.checkpoint:
+                # Tag all intermediate outputs as to be recomputed.
+                function_arg_names = {a.name for a in bsym.flat_proxy_args}
+
+                for subsym in bsym.subsymbols:
+                    for o in subsym.flat_proxy_outs:
+                        if o.name not in function_arg_names:
+                            o.tags.add(ProxyTag.RECOMPUTE_IN_BACKWARD)
+
+                # decompose
+                self.add_unprocessed_bsyms(bsym.subsymbols[:])
+                return
+
+            # 3a. see if we have a grad_transform (e.g. from OperatorExecutor.register_grad_transform)
+
+            # executor or global grad transform
+            joint_forward_backward, _ = _get_gradfn_and_executor(bsym)
+
+            if joint_forward_backward is None:
+                # 3b. if we don't have a grad_transform for a bsym, maybe we have old style augmented forward and backward
+                # registered. In this case, we make the joint_forward_backward which works like a grad_gransform.
+                # registered augmented forward impl (where aug fwd and backward are registered separately)
+                aug_fwd_impl = augmented_forward_impls.get(bsym.sym.id)
+                if aug_fwd_impl is not None:
+                    bwd_impl = backward_impls.get(bsym.sym.id)
+
+                    # this is our ad hoc combined forward and backward function ("grad_transform")
+                    def joint_forward_backward(*args, **kwargs):
+                        arg_proxy_names = {a.name for a in bsym.flat_proxy_args}
+
+                        # run the augmented forward res
+                        aug_fwd_res = aug_fwd_impl(*bsym.args, **bsym.kwargs)
+
+                        # aug_fwd_res could be either VJPDual or just a tuple, in any case, we can decompose it.
+                        res, saved_for_backward = aug_fwd_res
+
+                        # we need to shallow copy inputs that are returned for "get_grad" and "put_grad" to properly work
+                        # (this shallow copy is the equivalent because we of an "edge" in the PyTorch autograd graph)
+                        def shallow_copy_if_input(p):
+                            if isinstance(p, thunder.TensorProxy) and p.name in arg_proxy_names:
+                                return thunder.core.prims.shallow_copy(p)
+                            return p
+
+                        res = tree_map(shallow_copy_if_input, res)
+
+                        # now we need the backward. it starts by getting the grad_outs
+                        grad_outs = []
+                        for r in thunder.core.pytree.tree_iter(res):
+                            if isinstance(r, thunder.TensorProxy):
+                                grad_outs.append(prims.get_grad(r))
+
+                        # The backward computes the grad_inps of the bsym from the grad_outs
+                        # TODO: non-grad outputs of bwd?
+                        grad_inps = bwd_impl(*saved_for_backward, *grad_outs)
+                        if isinstance(grad_inps, thunder.Proxy):
+                            grad_inps = [grad_inps]
+
+                        # match the grad_inps to the inputs of the boudnd symbol and put the grads
+
+                        flat_inps = args
+                        # for autograd_function_apply, skip the function args
+                        # TODO: fix the returned gradients to include two None?.
+                        if bsym.sym == thunder.torch.autograd_function_apply:
+                            flat_inps = args[2:]
+
+                        # there may be non-gradient requiring additional args (todo: maybe only support this for non-tensor ones?)
+                        num_flat_tensor_inps = sum(isinstance(i, thunder.TensorProxy) for i in flat_inps)
+                        utils.check(
+                            num_flat_tensor_inps <= len(grad_inps),
+                            lambda: f"Backward for {bsym.sym.id} returned {len(grad_inps)} value(s), but expected {num_flat_tensor_inps}",
+                        )
+
+                        assert len(grad_inps) <= len(flat_inps)
+                        for i, gi in zip(flat_inps, grad_inps):
+                            # for integer proxies etc. we expect gi to be None
+                            if isinstance(i, thunder.TensorProxy) and gi is not None:
+                                prims.put_grad(i, gi)
+                        return res
+
+            # 3c if we have a joint forward backward (either grad_transform or the ad hoc constructed one above), we can comptue the gradient
+            if joint_forward_backward is not None:
+                self.new_trace.push_scope([])
+                result = joint_forward_backward(*bsym.args, **bsym.kwargs)
+
+                # Check that the number of outputs of the original forward function is the
+                # same as the number of primal outputs of the augmented forward trace
+                utils.check(
+                    len(utils.sequencify(bsym.output)) == len(utils.sequencify(result)),
+                    lambda: f"While generating forward and backward functions for {bsym.sym.name}, encountered an error.\n"
+                    "The number of outputs of the gradient transform function must be the same as the nunmber of outputs of the original forward function.\n"
+                    f"Number of outputs of the original forward function: {len(utils.sequencify(bsym.output))}\n"
+                    f"Number of primal outputs of the gradient tansform / augmented forward: {len(utils.sequencify(result))}\n"
path: thunder/transforms/autodiff.py
user: login: beverlylytle
body: ```suggestion
                    f"Number of primal outputs of the gradient transform / augmented forward: {len(utils.sequencify(result))}\n"
```
created_at: 2025-05-20T12:23:43Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,477 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims, utils
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+    recompute_saved_for_backward,
+)
+from thunder.core.proxies import ProxyTag
+from thunder.core.symbol import BoundSymbol
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+def _should_recompute_bsym_in_backward(bsym):
+    return any((ProxyTag.RECOMPUTE_IN_BACKWARD in o.tags) for o in bsym.flat_proxy_outs)
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                ########### BEGINNING of return handling (and putting the backward computation in the joint trace)
+                # This is big (and a bit messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols backward in order (from the generation of the list, this
+                # is the backward of the last forward instruction first, then then the backward of the second to last etc.
+                for bw_bsyms in self.collected_bw_part_bsyms:
+                    # if we don't have and grad_outs, we know that the grad_ins are not needed, so we don't compute them
+                    # (e.g. for `where(a * b > 0, a, b)` we do not need the grad components of a * b propagated
+                    if any(
+                        ((gg.name in grad_proxy_map) or (gg.name in output_proxy_names))
+                        for gg in (bs.args[0] for bs in bw_bsyms if bs.sym == prims.get_grad)
+                    ):
+                        for nbsym in bw_bsyms:
+                            if nbsym.sym == prims.put_grad:
+                                # put_grad adds something to the gradients
+                                # - if it is an output of the function, we also need to get the
+                                #   "grad_output" (but only once)
+                                # - if we already had a gradient, we need to add the new component
+                                # - structurally (for correctness), there should be no put_grad
+                                #   after a get_grad
+                                p = self.env.get(nbsym.args[0].name, nbsym.args[0])
+                                current_grad = grad_proxy_map.get(p.name)
+                                new_grad = nbsym.args[1]
+                                if current_grad is None and p.name in output_proxy_names:
+                                    self.new_trace.push_scope([])
+                                    current_grad = prims.get_grad(p)
+                                    self.add_processed_bsyms(self.new_trace.pop_scope())
+
+                                if current_grad is not None:
+                                    new_grad = self.add_bsyms_from_function(thunder.torch.add, current_grad, new_grad)
+
+                                grad_proxy_map[p.name] = new_grad
+                                self.write(new_grad, new_grad)
+                            elif nbsym.sym == prims.get_grad:
+                                # get grad reads a gradient, this could be
+                                # - the gradient of an intermediate (possibly combined with an output) form put_grad above
+                                #   note that all put_grads are happen before all get_grads from the dataflow of the forward,
+                                #   in this case we use the gradient from the put_grad(s)
+                                # - an output of the forward function (that has not been combined with an intermediate in pu grd),
+                                #   then the gradient is like the "grad_out" in the parameter list of a autograd.Function.backward
+                                #   in this case we keep a get_grad in the trace.
+                                # - an intermedite that does not have a gradient computed. In this case the gradient is 0 in the shape
+                                #   of the get_grad argument. This happens e.g. if we have something like "a > 0" for a float a then
+                                #   a has grad 0.
+                                #   TODO: We could (and should) optimize this (see below).
+                                p = nbsym.args[0]
+
+                                # TODO: set proxy requires grad here!
+                                name = p.name
+                                current_grad = grad_proxy_map.get(name)
+                                if current_grad is not None:
+                                    # do we also need to map?
+                                    self.write(nbsym.output, current_grad)
+                                    # replace_output_with_current_grad
+                                    self.swap_map[thunder.core.proxies.variableify(nbsym.output)] = current_grad
+                                elif name in output_proxy_names:
+                                    # output here???
+                                    new_bsym = nbsym.from_bsym()
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                else:
+                                    # TODO: mark this? if all inputs to the backward formula are unused, we would not want to compute it.
+                                    new_bsym = thunder.torch.zeros.bind(
+                                        *p.shape, device=p.device, dtype=p.dtype, output=nbsym.output
+                                    )
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                    self.write(nbsym.output, new_bsym.output)
+                            else:
+                                # all other symbols are just computation
+                                self.add_processed_bsyms([nbsym.from_bsym()])
+                self.collected_bw_part_bsyms.clear()
+
+                # we collect the gradients fo the flat args in the output dictionary's 'grad_flat_args'
+                # this means that the joint trace as them as outputs, which is good.
+                grad_flat_args = []
+                for p in bsym.args[0]["flat_args"]:
+                    # or p = self.read(p) here?
+                    if isinstance(p, thunder.TensorProxy) and p.requires_grad and p.name in grad_proxy_map:
+                        # is it always OK if we don't have a gradient? (one case: unused input)
+                        # result of put_grad???
+                        grad_flat_args.append(grad_proxy_map[p.name])
+                    else:
+                        grad_flat_args.append(None)
+
+                bsym.args[0]["grad_flat_args"] = grad_flat_args
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+                ########### END of return handling (and putting the backward computation in the joint trace)
+
+            # we now handle various bound symbols by collecting augmented forward and backward symbols
+
+            # 1. constant for gradients (no grad required)
+            #    excemting synchronize here terrible hack to cope with non-grad-needing sharded tensors
+            #    as required by LoRA. We should have a symbol tag "always compute grad" instead.
+            if is_constant_for_vjp(bsym) and not bsym.sym.name == "synchronize":
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+
+            # 2. Special case the thunder.torch.checkpoint higher oder function
+            if bsym.sym == thunder.torch.checkpoint:
+                # Tag all intermediate outputs as to be recomputed.
+                function_arg_names = {a.name for a in bsym.flat_proxy_args}
+
+                for subsym in bsym.subsymbols:
+                    for o in subsym.flat_proxy_outs:
+                        if o.name not in function_arg_names:
+                            o.tags.add(ProxyTag.RECOMPUTE_IN_BACKWARD)
+
+                # decompose
+                self.add_unprocessed_bsyms(bsym.subsymbols[:])
+                return
+
+            # 3a. see if we have a grad_transform (e.g. from OperatorExecutor.register_grad_transform)
+
+            # executor or global grad transform
+            joint_forward_backward, _ = _get_gradfn_and_executor(bsym)
+
+            if joint_forward_backward is None:
+                # 3b. if we don't have a grad_transform for a bsym, maybe we have old style augmented forward and backward
+                # registered. In this case, we make the joint_forward_backward which works like a grad_gransform.
+                # registered augmented forward impl (where aug fwd and backward are registered separately)
+                aug_fwd_impl = augmented_forward_impls.get(bsym.sym.id)
+                if aug_fwd_impl is not None:
+                    bwd_impl = backward_impls.get(bsym.sym.id)
+
+                    # this is our ad hoc combined forward and backward function ("grad_transform")
+                    def joint_forward_backward(*args, **kwargs):
+                        arg_proxy_names = {a.name for a in bsym.flat_proxy_args}
+
+                        # run the augmented forward res
+                        aug_fwd_res = aug_fwd_impl(*bsym.args, **bsym.kwargs)
+
+                        # aug_fwd_res could be either VJPDual or just a tuple, in any case, we can decompose it.
+                        res, saved_for_backward = aug_fwd_res
+
+                        # we need to shallow copy inputs that are returned for "get_grad" and "put_grad" to properly work
+                        # (this shallow copy is the equivalent because we of an "edge" in the PyTorch autograd graph)
+                        def shallow_copy_if_input(p):
+                            if isinstance(p, thunder.TensorProxy) and p.name in arg_proxy_names:
+                                return thunder.core.prims.shallow_copy(p)
+                            return p
+
+                        res = tree_map(shallow_copy_if_input, res)
+
+                        # now we need the backward. it starts by getting the grad_outs
+                        grad_outs = []
+                        for r in thunder.core.pytree.tree_iter(res):
+                            if isinstance(r, thunder.TensorProxy):
+                                grad_outs.append(prims.get_grad(r))
+
+                        # The backward computes the grad_inps of the bsym from the grad_outs
+                        # TODO: non-grad outputs of bwd?
+                        grad_inps = bwd_impl(*saved_for_backward, *grad_outs)
+                        if isinstance(grad_inps, thunder.Proxy):
+                            grad_inps = [grad_inps]
+
+                        # match the grad_inps to the inputs of the boudnd symbol and put the grads
+
+                        flat_inps = args
+                        # for autograd_function_apply, skip the function args
+                        # TODO: fix the returned gradients to include two None?.
+                        if bsym.sym == thunder.torch.autograd_function_apply:
+                            flat_inps = args[2:]
+
+                        # there may be non-gradient requiring additional args (todo: maybe only support this for non-tensor ones?)
+                        num_flat_tensor_inps = sum(isinstance(i, thunder.TensorProxy) for i in flat_inps)
+                        utils.check(
+                            num_flat_tensor_inps <= len(grad_inps),
+                            lambda: f"Backward for {bsym.sym.id} returned {len(grad_inps)} value(s), but expected {num_flat_tensor_inps}",
+                        )
+
+                        assert len(grad_inps) <= len(flat_inps)
+                        for i, gi in zip(flat_inps, grad_inps):
+                            # for integer proxies etc. we expect gi to be None
+                            if isinstance(i, thunder.TensorProxy) and gi is not None:
+                                prims.put_grad(i, gi)
+                        return res
+
+            # 3c if we have a joint forward backward (either grad_transform or the ad hoc constructed one above), we can comptue the gradient
+            if joint_forward_backward is not None:
+                self.new_trace.push_scope([])
+                result = joint_forward_backward(*bsym.args, **bsym.kwargs)
+
+                # Check that the number of outputs of the original forward function is the
+                # same as the number of primal outputs of the augmented forward trace
+                utils.check(
+                    len(utils.sequencify(bsym.output)) == len(utils.sequencify(result)),
+                    lambda: f"While generating forward and backward functions for {bsym.sym.name}, encountered an error.\n"
+                    "The number of outputs of the gradient transform function must be the same as the nunmber of outputs of the original forward function.\n"
+                    f"Number of outputs of the original forward function: {len(utils.sequencify(bsym.output))}\n"
+                    f"Number of primal outputs of the gradient tansform / augmented forward: {len(utils.sequencify(result))}\n"
+                    "Please check the forward function and the gradient transfrm function / augmented forward to ensure that they have the same number of outputs.",
+                )
+                self.set_result(result)
+                new_bsyms = self.new_trace.pop_scope()
+
+                # simple splitting: only compute in forward what is needed for the output
+                forward_part_proxy_names = {
+                    o.name for o in thunder.core.pytree.tree_iter(result) if isinstance(o, thunder.Proxy)
+                }
+                forward_part_bsyms = []
+                backward_part_bsyms = []
+                for nbsym in reversed(new_bsyms):
+                    # argnames = {p.name for p in nbsym.flat_proxy_args}
+                    if nbsym.sym in {prims.get_grad, prims.put_grad}:
+                        backward_part_bsyms.insert(0, nbsym)
+                        continue
+                    assert nbsym.sym != prims.unpack_trivial  # can unpack trivial happen here?
+                    if any((o.name in forward_part_proxy_names) for o in nbsym.flat_proxy_outs):
+                        forward_part_bsyms.insert(0, nbsym)
+                        bsym_arg_names = {a.name for a in nbsym.flat_proxy_args}
+                        bsym_out_names = {o.name for o in nbsym.flat_proxy_outs if o.name not in bsym_arg_names}
+                        forward_part_proxy_names.update(bsym_arg_names)
+                        forward_part_proxy_names.update(bsym_out_names)
+                        continue
+                    backward_part_bsyms.insert(0, nbsym)
+                self.add_processed_bsyms(forward_part_bsyms)
+
+                self.collected_bw_part_bsyms.insert(0, backward_part_bsyms)
+                return
+
+            # 4. No gradient transform found, need to decompose the symbol
+
+            # error if this is a primitive
+            utils.check(not bsym.sym.is_prim, lambda: f"Failed to find a gradient transform for bound symbol {bsym=}")
+
+            # TODO: check if this is needed: the old impl checked whether len(bsym.subsymbols) > 0 except for the special case "torch.nn.functional.dropout" with p=0...
+            # add the decomposition (= the subsymbols) to the front of the symbols to be processed
+            self.add_unprocessed_bsyms(bsym.subsymbols[:])
+
+            # end of 4 and end of the bsym processing loop.
+
+    start_time_ns = time.perf_counter_ns()
+
+    # run the trace through the processor
+    trace, _ = AugmentedForwardProcessor(trace)()
+    # run through DCE in case some of the gradients of intermediates are not needed.
+    trace = thunder.core.transform_common.dce(trace)
+
+    end_time_ns = time.perf_counter_ns()
+    elapsed_time_ns = end_time_ns - start_time_ns
+    elapsed_time_millis = elapsed_time_ns // 1000000
+    trace.set_provenance(
+        thunder.core.trace.TraceProvenance(f"Grad transform pass (took {elapsed_time_millis} milliseconds)")
+    )
+    return trace
+
+
+def split_into_forward_and_backward(joint_trace):
+    """split a joint trace for forward and backward into separate ones, including recomputation (aka activation checkpointing)"""
+
+    # the joint trace will have the forward computation at the beginning and then the backward computation
+    # from how it is constructed.
+    # we split the trace:
+    # - forward symbols go into forward_part_bsyms
+    # - all symbols not in the forward go into backward_part_bsyms
+    # - for recomputation (aka activation checkpointing), we want to insert symbols going into the forward also into the
+    #   backward. but we want to do so "just in time". To this end, we gather the symbols in a dict and later
+    #   insert it when their respective outputs are needed. This is in backward_part_bsyms_recomputated
+    #   The just in time recomputation is a heuristic to save memory mimicking checkpointing: e.g. for a checkpointed
+    #   block, the forward would be recomputed just before computing the gradient.
+    # the splitting is done in the reverse order fo the bound symbols, and works out which bits are needed for the forward
path: thunder/transforms/autodiff.py
user: login: beverlylytle
body: ```suggestion
    # the splitting is done in the reverse order of the bound symbols, and works out which bits are needed for the forward
```
created_at: 2025-05-20T12:25:41Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,477 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims, utils
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+    recompute_saved_for_backward,
+)
+from thunder.core.proxies import ProxyTag
+from thunder.core.symbol import BoundSymbol
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+def _should_recompute_bsym_in_backward(bsym):
+    return any((ProxyTag.RECOMPUTE_IN_BACKWARD in o.tags) for o in bsym.flat_proxy_outs)
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                ########### BEGINNING of return handling (and putting the backward computation in the joint trace)
+                # This is big (and a bit messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols backward in order (from the generation of the list, this
+                # is the backward of the last forward instruction first, then then the backward of the second to last etc.
+                for bw_bsyms in self.collected_bw_part_bsyms:
+                    # if we don't have and grad_outs, we know that the grad_ins are not needed, so we don't compute them
+                    # (e.g. for `where(a * b > 0, a, b)` we do not need the grad components of a * b propagated
+                    if any(
+                        ((gg.name in grad_proxy_map) or (gg.name in output_proxy_names))
+                        for gg in (bs.args[0] for bs in bw_bsyms if bs.sym == prims.get_grad)
+                    ):
+                        for nbsym in bw_bsyms:
+                            if nbsym.sym == prims.put_grad:
+                                # put_grad adds something to the gradients
+                                # - if it is an output of the function, we also need to get the
+                                #   "grad_output" (but only once)
+                                # - if we already had a gradient, we need to add the new component
+                                # - structurally (for correctness), there should be no put_grad
+                                #   after a get_grad
+                                p = self.env.get(nbsym.args[0].name, nbsym.args[0])
+                                current_grad = grad_proxy_map.get(p.name)
+                                new_grad = nbsym.args[1]
+                                if current_grad is None and p.name in output_proxy_names:
+                                    self.new_trace.push_scope([])
+                                    current_grad = prims.get_grad(p)
+                                    self.add_processed_bsyms(self.new_trace.pop_scope())
+
+                                if current_grad is not None:
+                                    new_grad = self.add_bsyms_from_function(thunder.torch.add, current_grad, new_grad)
+
+                                grad_proxy_map[p.name] = new_grad
+                                self.write(new_grad, new_grad)
+                            elif nbsym.sym == prims.get_grad:
+                                # get grad reads a gradient, this could be
+                                # - the gradient of an intermediate (possibly combined with an output) form put_grad above
+                                #   note that all put_grads are happen before all get_grads from the dataflow of the forward,
+                                #   in this case we use the gradient from the put_grad(s)
+                                # - an output of the forward function (that has not been combined with an intermediate in pu grd),
+                                #   then the gradient is like the "grad_out" in the parameter list of a autograd.Function.backward
+                                #   in this case we keep a get_grad in the trace.
+                                # - an intermedite that does not have a gradient computed. In this case the gradient is 0 in the shape
+                                #   of the get_grad argument. This happens e.g. if we have something like "a > 0" for a float a then
+                                #   a has grad 0.
+                                #   TODO: We could (and should) optimize this (see below).
+                                p = nbsym.args[0]
+
+                                # TODO: set proxy requires grad here!
+                                name = p.name
+                                current_grad = grad_proxy_map.get(name)
+                                if current_grad is not None:
+                                    # do we also need to map?
+                                    self.write(nbsym.output, current_grad)
+                                    # replace_output_with_current_grad
+                                    self.swap_map[thunder.core.proxies.variableify(nbsym.output)] = current_grad
+                                elif name in output_proxy_names:
+                                    # output here???
+                                    new_bsym = nbsym.from_bsym()
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                else:
+                                    # TODO: mark this? if all inputs to the backward formula are unused, we would not want to compute it.
+                                    new_bsym = thunder.torch.zeros.bind(
+                                        *p.shape, device=p.device, dtype=p.dtype, output=nbsym.output
+                                    )
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                    self.write(nbsym.output, new_bsym.output)
+                            else:
+                                # all other symbols are just computation
+                                self.add_processed_bsyms([nbsym.from_bsym()])
+                self.collected_bw_part_bsyms.clear()
+
+                # we collect the gradients fo the flat args in the output dictionary's 'grad_flat_args'
+                # this means that the joint trace as them as outputs, which is good.
+                grad_flat_args = []
+                for p in bsym.args[0]["flat_args"]:
+                    # or p = self.read(p) here?
+                    if isinstance(p, thunder.TensorProxy) and p.requires_grad and p.name in grad_proxy_map:
+                        # is it always OK if we don't have a gradient? (one case: unused input)
+                        # result of put_grad???
+                        grad_flat_args.append(grad_proxy_map[p.name])
+                    else:
+                        grad_flat_args.append(None)
+
+                bsym.args[0]["grad_flat_args"] = grad_flat_args
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+                ########### END of return handling (and putting the backward computation in the joint trace)
+
+            # we now handle various bound symbols by collecting augmented forward and backward symbols
+
+            # 1. constant for gradients (no grad required)
+            #    excemting synchronize here terrible hack to cope with non-grad-needing sharded tensors
+            #    as required by LoRA. We should have a symbol tag "always compute grad" instead.
+            if is_constant_for_vjp(bsym) and not bsym.sym.name == "synchronize":
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+
+            # 2. Special case the thunder.torch.checkpoint higher oder function
+            if bsym.sym == thunder.torch.checkpoint:
+                # Tag all intermediate outputs as to be recomputed.
+                function_arg_names = {a.name for a in bsym.flat_proxy_args}
+
+                for subsym in bsym.subsymbols:
+                    for o in subsym.flat_proxy_outs:
+                        if o.name not in function_arg_names:
+                            o.tags.add(ProxyTag.RECOMPUTE_IN_BACKWARD)
+
+                # decompose
+                self.add_unprocessed_bsyms(bsym.subsymbols[:])
+                return
+
+            # 3a. see if we have a grad_transform (e.g. from OperatorExecutor.register_grad_transform)
+
+            # executor or global grad transform
+            joint_forward_backward, _ = _get_gradfn_and_executor(bsym)
+
+            if joint_forward_backward is None:
+                # 3b. if we don't have a grad_transform for a bsym, maybe we have old style augmented forward and backward
+                # registered. In this case, we make the joint_forward_backward which works like a grad_gransform.
+                # registered augmented forward impl (where aug fwd and backward are registered separately)
+                aug_fwd_impl = augmented_forward_impls.get(bsym.sym.id)
+                if aug_fwd_impl is not None:
+                    bwd_impl = backward_impls.get(bsym.sym.id)
+
+                    # this is our ad hoc combined forward and backward function ("grad_transform")
+                    def joint_forward_backward(*args, **kwargs):
+                        arg_proxy_names = {a.name for a in bsym.flat_proxy_args}
+
+                        # run the augmented forward res
+                        aug_fwd_res = aug_fwd_impl(*bsym.args, **bsym.kwargs)
+
+                        # aug_fwd_res could be either VJPDual or just a tuple, in any case, we can decompose it.
+                        res, saved_for_backward = aug_fwd_res
+
+                        # we need to shallow copy inputs that are returned for "get_grad" and "put_grad" to properly work
+                        # (this shallow copy is the equivalent because we of an "edge" in the PyTorch autograd graph)
+                        def shallow_copy_if_input(p):
+                            if isinstance(p, thunder.TensorProxy) and p.name in arg_proxy_names:
+                                return thunder.core.prims.shallow_copy(p)
+                            return p
+
+                        res = tree_map(shallow_copy_if_input, res)
+
+                        # now we need the backward. it starts by getting the grad_outs
+                        grad_outs = []
+                        for r in thunder.core.pytree.tree_iter(res):
+                            if isinstance(r, thunder.TensorProxy):
+                                grad_outs.append(prims.get_grad(r))
+
+                        # The backward computes the grad_inps of the bsym from the grad_outs
+                        # TODO: non-grad outputs of bwd?
+                        grad_inps = bwd_impl(*saved_for_backward, *grad_outs)
+                        if isinstance(grad_inps, thunder.Proxy):
+                            grad_inps = [grad_inps]
+
+                        # match the grad_inps to the inputs of the boudnd symbol and put the grads
+
+                        flat_inps = args
+                        # for autograd_function_apply, skip the function args
+                        # TODO: fix the returned gradients to include two None?.
+                        if bsym.sym == thunder.torch.autograd_function_apply:
+                            flat_inps = args[2:]
+
+                        # there may be non-gradient requiring additional args (todo: maybe only support this for non-tensor ones?)
+                        num_flat_tensor_inps = sum(isinstance(i, thunder.TensorProxy) for i in flat_inps)
+                        utils.check(
+                            num_flat_tensor_inps <= len(grad_inps),
+                            lambda: f"Backward for {bsym.sym.id} returned {len(grad_inps)} value(s), but expected {num_flat_tensor_inps}",
+                        )
+
+                        assert len(grad_inps) <= len(flat_inps)
+                        for i, gi in zip(flat_inps, grad_inps):
+                            # for integer proxies etc. we expect gi to be None
+                            if isinstance(i, thunder.TensorProxy) and gi is not None:
+                                prims.put_grad(i, gi)
+                        return res
+
+            # 3c if we have a joint forward backward (either grad_transform or the ad hoc constructed one above), we can comptue the gradient
+            if joint_forward_backward is not None:
+                self.new_trace.push_scope([])
+                result = joint_forward_backward(*bsym.args, **bsym.kwargs)
+
+                # Check that the number of outputs of the original forward function is the
+                # same as the number of primal outputs of the augmented forward trace
+                utils.check(
+                    len(utils.sequencify(bsym.output)) == len(utils.sequencify(result)),
+                    lambda: f"While generating forward and backward functions for {bsym.sym.name}, encountered an error.\n"
+                    "The number of outputs of the gradient transform function must be the same as the nunmber of outputs of the original forward function.\n"
+                    f"Number of outputs of the original forward function: {len(utils.sequencify(bsym.output))}\n"
+                    f"Number of primal outputs of the gradient tansform / augmented forward: {len(utils.sequencify(result))}\n"
+                    "Please check the forward function and the gradient transfrm function / augmented forward to ensure that they have the same number of outputs.",
+                )
+                self.set_result(result)
+                new_bsyms = self.new_trace.pop_scope()
+
+                # simple splitting: only compute in forward what is needed for the output
+                forward_part_proxy_names = {
+                    o.name for o in thunder.core.pytree.tree_iter(result) if isinstance(o, thunder.Proxy)
+                }
+                forward_part_bsyms = []
+                backward_part_bsyms = []
+                for nbsym in reversed(new_bsyms):
+                    # argnames = {p.name for p in nbsym.flat_proxy_args}
+                    if nbsym.sym in {prims.get_grad, prims.put_grad}:
+                        backward_part_bsyms.insert(0, nbsym)
+                        continue
+                    assert nbsym.sym != prims.unpack_trivial  # can unpack trivial happen here?
+                    if any((o.name in forward_part_proxy_names) for o in nbsym.flat_proxy_outs):
+                        forward_part_bsyms.insert(0, nbsym)
+                        bsym_arg_names = {a.name for a in nbsym.flat_proxy_args}
+                        bsym_out_names = {o.name for o in nbsym.flat_proxy_outs if o.name not in bsym_arg_names}
+                        forward_part_proxy_names.update(bsym_arg_names)
+                        forward_part_proxy_names.update(bsym_out_names)
+                        continue
+                    backward_part_bsyms.insert(0, nbsym)
+                self.add_processed_bsyms(forward_part_bsyms)
+
+                self.collected_bw_part_bsyms.insert(0, backward_part_bsyms)
+                return
+
+            # 4. No gradient transform found, need to decompose the symbol
+
+            # error if this is a primitive
+            utils.check(not bsym.sym.is_prim, lambda: f"Failed to find a gradient transform for bound symbol {bsym=}")
+
+            # TODO: check if this is needed: the old impl checked whether len(bsym.subsymbols) > 0 except for the special case "torch.nn.functional.dropout" with p=0...
+            # add the decomposition (= the subsymbols) to the front of the symbols to be processed
+            self.add_unprocessed_bsyms(bsym.subsymbols[:])
+
+            # end of 4 and end of the bsym processing loop.
+
+    start_time_ns = time.perf_counter_ns()
+
+    # run the trace through the processor
+    trace, _ = AugmentedForwardProcessor(trace)()
+    # run through DCE in case some of the gradients of intermediates are not needed.
+    trace = thunder.core.transform_common.dce(trace)
+
+    end_time_ns = time.perf_counter_ns()
+    elapsed_time_ns = end_time_ns - start_time_ns
+    elapsed_time_millis = elapsed_time_ns // 1000000
+    trace.set_provenance(
+        thunder.core.trace.TraceProvenance(f"Grad transform pass (took {elapsed_time_millis} milliseconds)")
+    )
+    return trace
+
+
+def split_into_forward_and_backward(joint_trace):
+    """split a joint trace for forward and backward into separate ones, including recomputation (aka activation checkpointing)"""
+
+    # the joint trace will have the forward computation at the beginning and then the backward computation
+    # from how it is constructed.
+    # we split the trace:
+    # - forward symbols go into forward_part_bsyms
+    # - all symbols not in the forward go into backward_part_bsyms
+    # - for recomputation (aka activation checkpointing), we want to insert symbols going into the forward also into the
+    #   backward. but we want to do so "just in time". To this end, we gather the symbols in a dict and later
+    #   insert it when their respective outputs are needed. This is in backward_part_bsyms_recomputated
+    #   The just in time recomputation is a heuristic to save memory mimicking checkpointing: e.g. for a checkpointed
+    #   block, the forward would be recomputed just before computing the gradient.
+    # the splitting is done in the reverse order fo the bound symbols, and works out which bits are needed for the forward
+    # from there.
+    forward_part_bsyms = []
+    backward_part_bsyms = []
+    backward_part_bsyms_recomputed: dict[str, tuple[BoundSymbol, set[str]]] = {}  # name -> (bsym, name_set)
+
+    # return has a dict, it contains the forward outputs in "output" and the backward outputs in "grad_flat_args", corresponding
+    # to "flat_args"
+    return_bsym = joint_trace.bound_symbols[-1]
+    assert return_bsym.sym == prims.python_return
+    fw_output = return_bsym.args[0]["output"]
+    assert isinstance(fw_output, tuple)
+
+    grad_outs = [None for _ in fw_output]
+    output_pos = {o.name: i for i, o in enumerate(fw_output) if isinstance(o, thunder.TensorProxy)}
+
+    # the proxies we need to compute in the forward - we start with the outputs of the forward
+    forward_proxy_names = {o.name for o in thunder.core.pytree.tree_iter(fw_output) if isinstance(o, thunder.Proxy)}
+    # we also have the inputs available, so we add flat_args.
+    # for inplace, we need to update this (or have flat args be the right thing?...)
+    forward_proxy_names.update(a.name for a in return_bsym.args[0]["flat_args"] if isinstance(a, thunder.Proxy))
+
+    # We keep track of the names of proxies we recompute in the backward as those will not need to be part of the
+    # ones saved in the forward for the backward
+    backward_recomputed_proxy_names = set()
+
+    # here we go through the bound symbols to assign to the three categories, keeping track of the forward names, the backward needed
+    # and the backward available names.
path: thunder/transforms/autodiff.py
user: login: beverlylytle
body: I think this comment is stale, as there is now only `backward_recomputed_proxy_names`.
created_at: 2025-05-20T12:26:56Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,477 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims, utils
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+    recompute_saved_for_backward,
+)
+from thunder.core.proxies import ProxyTag
+from thunder.core.symbol import BoundSymbol
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+def _should_recompute_bsym_in_backward(bsym):
+    return any((ProxyTag.RECOMPUTE_IN_BACKWARD in o.tags) for o in bsym.flat_proxy_outs)
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                ########### BEGINNING of return handling (and putting the backward computation in the joint trace)
+                # This is big (and a bit messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols backward in order (from the generation of the list, this
+                # is the backward of the last forward instruction first, then then the backward of the second to last etc.
+                for bw_bsyms in self.collected_bw_part_bsyms:
+                    # if we don't have and grad_outs, we know that the grad_ins are not needed, so we don't compute them
+                    # (e.g. for `where(a * b > 0, a, b)` we do not need the grad components of a * b propagated
+                    if any(
+                        ((gg.name in grad_proxy_map) or (gg.name in output_proxy_names))
+                        for gg in (bs.args[0] for bs in bw_bsyms if bs.sym == prims.get_grad)
+                    ):
+                        for nbsym in bw_bsyms:
+                            if nbsym.sym == prims.put_grad:
+                                # put_grad adds something to the gradients
+                                # - if it is an output of the function, we also need to get the
+                                #   "grad_output" (but only once)
+                                # - if we already had a gradient, we need to add the new component
+                                # - structurally (for correctness), there should be no put_grad
+                                #   after a get_grad
+                                p = self.env.get(nbsym.args[0].name, nbsym.args[0])
+                                current_grad = grad_proxy_map.get(p.name)
+                                new_grad = nbsym.args[1]
+                                if current_grad is None and p.name in output_proxy_names:
+                                    self.new_trace.push_scope([])
+                                    current_grad = prims.get_grad(p)
+                                    self.add_processed_bsyms(self.new_trace.pop_scope())
+
+                                if current_grad is not None:
+                                    new_grad = self.add_bsyms_from_function(thunder.torch.add, current_grad, new_grad)
+
+                                grad_proxy_map[p.name] = new_grad
+                                self.write(new_grad, new_grad)
+                            elif nbsym.sym == prims.get_grad:
+                                # get grad reads a gradient, this could be
+                                # - the gradient of an intermediate (possibly combined with an output) form put_grad above
+                                #   note that all put_grads are happen before all get_grads from the dataflow of the forward,
+                                #   in this case we use the gradient from the put_grad(s)
+                                # - an output of the forward function (that has not been combined with an intermediate in pu grd),
+                                #   then the gradient is like the "grad_out" in the parameter list of a autograd.Function.backward
+                                #   in this case we keep a get_grad in the trace.
+                                # - an intermedite that does not have a gradient computed. In this case the gradient is 0 in the shape
+                                #   of the get_grad argument. This happens e.g. if we have something like "a > 0" for a float a then
+                                #   a has grad 0.
+                                #   TODO: We could (and should) optimize this (see below).
+                                p = nbsym.args[0]
+
+                                # TODO: set proxy requires grad here!
+                                name = p.name
+                                current_grad = grad_proxy_map.get(name)
+                                if current_grad is not None:
+                                    # do we also need to map?
+                                    self.write(nbsym.output, current_grad)
+                                    # replace_output_with_current_grad
+                                    self.swap_map[thunder.core.proxies.variableify(nbsym.output)] = current_grad
+                                elif name in output_proxy_names:
+                                    # output here???
+                                    new_bsym = nbsym.from_bsym()
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                else:
+                                    # TODO: mark this? if all inputs to the backward formula are unused, we would not want to compute it.
+                                    new_bsym = thunder.torch.zeros.bind(
+                                        *p.shape, device=p.device, dtype=p.dtype, output=nbsym.output
+                                    )
+                                    self.add_processed_bsyms([new_bsym])
+                                    grad_proxy_map[name] = new_bsym.output
+                                    self.write(nbsym.output, new_bsym.output)
+                            else:
+                                # all other symbols are just computation
+                                self.add_processed_bsyms([nbsym.from_bsym()])
+                self.collected_bw_part_bsyms.clear()
+
+                # we collect the gradients fo the flat args in the output dictionary's 'grad_flat_args'
+                # this means that the joint trace as them as outputs, which is good.
+                grad_flat_args = []
+                for p in bsym.args[0]["flat_args"]:
+                    # or p = self.read(p) here?
+                    if isinstance(p, thunder.TensorProxy) and p.requires_grad and p.name in grad_proxy_map:
+                        # is it always OK if we don't have a gradient? (one case: unused input)
+                        # result of put_grad???
+                        grad_flat_args.append(grad_proxy_map[p.name])
+                    else:
+                        grad_flat_args.append(None)
+
+                bsym.args[0]["grad_flat_args"] = grad_flat_args
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+                ########### END of return handling (and putting the backward computation in the joint trace)
+
+            # we now handle various bound symbols by collecting augmented forward and backward symbols
+
+            # 1. constant for gradients (no grad required)
+            #    excemting synchronize here terrible hack to cope with non-grad-needing sharded tensors
+            #    as required by LoRA. We should have a symbol tag "always compute grad" instead.
+            if is_constant_for_vjp(bsym) and not bsym.sym.name == "synchronize":
+                self.add_processed_bsyms([bsym.from_bsym()])
+                self.set_result(bsym.output)
+                return
+
+            # 2. Special case the thunder.torch.checkpoint higher oder function
+            if bsym.sym == thunder.torch.checkpoint:
+                # Tag all intermediate outputs as to be recomputed.
+                function_arg_names = {a.name for a in bsym.flat_proxy_args}
+
+                for subsym in bsym.subsymbols:
+                    for o in subsym.flat_proxy_outs:
+                        if o.name not in function_arg_names:
+                            o.tags.add(ProxyTag.RECOMPUTE_IN_BACKWARD)
+
+                # decompose
+                self.add_unprocessed_bsyms(bsym.subsymbols[:])
+                return
+
+            # 3a. see if we have a grad_transform (e.g. from OperatorExecutor.register_grad_transform)
+
+            # executor or global grad transform
+            joint_forward_backward, _ = _get_gradfn_and_executor(bsym)
+
+            if joint_forward_backward is None:
+                # 3b. if we don't have a grad_transform for a bsym, maybe we have old style augmented forward and backward
+                # registered. In this case, we make the joint_forward_backward which works like a grad_gransform.
+                # registered augmented forward impl (where aug fwd and backward are registered separately)
+                aug_fwd_impl = augmented_forward_impls.get(bsym.sym.id)
+                if aug_fwd_impl is not None:
+                    bwd_impl = backward_impls.get(bsym.sym.id)
+
+                    # this is our ad hoc combined forward and backward function ("grad_transform")
+                    def joint_forward_backward(*args, **kwargs):
+                        arg_proxy_names = {a.name for a in bsym.flat_proxy_args}
+
+                        # run the augmented forward res
+                        aug_fwd_res = aug_fwd_impl(*bsym.args, **bsym.kwargs)
+
+                        # aug_fwd_res could be either VJPDual or just a tuple, in any case, we can decompose it.
+                        res, saved_for_backward = aug_fwd_res
+
+                        # we need to shallow copy inputs that are returned for "get_grad" and "put_grad" to properly work
+                        # (this shallow copy is the equivalent because we of an "edge" in the PyTorch autograd graph)
+                        def shallow_copy_if_input(p):
+                            if isinstance(p, thunder.TensorProxy) and p.name in arg_proxy_names:
+                                return thunder.core.prims.shallow_copy(p)
+                            return p
+
+                        res = tree_map(shallow_copy_if_input, res)
+
+                        # now we need the backward. it starts by getting the grad_outs
+                        grad_outs = []
+                        for r in thunder.core.pytree.tree_iter(res):
+                            if isinstance(r, thunder.TensorProxy):
+                                grad_outs.append(prims.get_grad(r))
+
+                        # The backward computes the grad_inps of the bsym from the grad_outs
+                        # TODO: non-grad outputs of bwd?
+                        grad_inps = bwd_impl(*saved_for_backward, *grad_outs)
+                        if isinstance(grad_inps, thunder.Proxy):
+                            grad_inps = [grad_inps]
+
+                        # match the grad_inps to the inputs of the boudnd symbol and put the grads
+
+                        flat_inps = args
+                        # for autograd_function_apply, skip the function args
+                        # TODO: fix the returned gradients to include two None?.
+                        if bsym.sym == thunder.torch.autograd_function_apply:
+                            flat_inps = args[2:]
+
+                        # there may be non-gradient requiring additional args (todo: maybe only support this for non-tensor ones?)
+                        num_flat_tensor_inps = sum(isinstance(i, thunder.TensorProxy) for i in flat_inps)
+                        utils.check(
+                            num_flat_tensor_inps <= len(grad_inps),
+                            lambda: f"Backward for {bsym.sym.id} returned {len(grad_inps)} value(s), but expected {num_flat_tensor_inps}",
+                        )
+
+                        assert len(grad_inps) <= len(flat_inps)
+                        for i, gi in zip(flat_inps, grad_inps):
+                            # for integer proxies etc. we expect gi to be None
+                            if isinstance(i, thunder.TensorProxy) and gi is not None:
+                                prims.put_grad(i, gi)
+                        return res
+
+            # 3c if we have a joint forward backward (either grad_transform or the ad hoc constructed one above), we can comptue the gradient
+            if joint_forward_backward is not None:
+                self.new_trace.push_scope([])
+                result = joint_forward_backward(*bsym.args, **bsym.kwargs)
+
+                # Check that the number of outputs of the original forward function is the
+                # same as the number of primal outputs of the augmented forward trace
+                utils.check(
+                    len(utils.sequencify(bsym.output)) == len(utils.sequencify(result)),
+                    lambda: f"While generating forward and backward functions for {bsym.sym.name}, encountered an error.\n"
+                    "The number of outputs of the gradient transform function must be the same as the nunmber of outputs of the original forward function.\n"
+                    f"Number of outputs of the original forward function: {len(utils.sequencify(bsym.output))}\n"
+                    f"Number of primal outputs of the gradient tansform / augmented forward: {len(utils.sequencify(result))}\n"
+                    "Please check the forward function and the gradient transfrm function / augmented forward to ensure that they have the same number of outputs.",
+                )
+                self.set_result(result)
+                new_bsyms = self.new_trace.pop_scope()
+
+                # simple splitting: only compute in forward what is needed for the output
+                forward_part_proxy_names = {
+                    o.name for o in thunder.core.pytree.tree_iter(result) if isinstance(o, thunder.Proxy)
+                }
+                forward_part_bsyms = []
+                backward_part_bsyms = []
+                for nbsym in reversed(new_bsyms):
+                    # argnames = {p.name for p in nbsym.flat_proxy_args}
+                    if nbsym.sym in {prims.get_grad, prims.put_grad}:
+                        backward_part_bsyms.insert(0, nbsym)
+                        continue
+                    assert nbsym.sym != prims.unpack_trivial  # can unpack trivial happen here?
+                    if any((o.name in forward_part_proxy_names) for o in nbsym.flat_proxy_outs):
+                        forward_part_bsyms.insert(0, nbsym)
+                        bsym_arg_names = {a.name for a in nbsym.flat_proxy_args}
+                        bsym_out_names = {o.name for o in nbsym.flat_proxy_outs if o.name not in bsym_arg_names}
+                        forward_part_proxy_names.update(bsym_arg_names)
+                        forward_part_proxy_names.update(bsym_out_names)
+                        continue
+                    backward_part_bsyms.insert(0, nbsym)
+                self.add_processed_bsyms(forward_part_bsyms)
+
+                self.collected_bw_part_bsyms.insert(0, backward_part_bsyms)
+                return
+
+            # 4. No gradient transform found, need to decompose the symbol
+
+            # error if this is a primitive
+            utils.check(not bsym.sym.is_prim, lambda: f"Failed to find a gradient transform for bound symbol {bsym=}")
+
+            # TODO: check if this is needed: the old impl checked whether len(bsym.subsymbols) > 0 except for the special case "torch.nn.functional.dropout" with p=0...
+            # add the decomposition (= the subsymbols) to the front of the symbols to be processed
+            self.add_unprocessed_bsyms(bsym.subsymbols[:])
+
+            # end of 4 and end of the bsym processing loop.
+
+    start_time_ns = time.perf_counter_ns()
+
+    # run the trace through the processor
+    trace, _ = AugmentedForwardProcessor(trace)()
+    # run through DCE in case some of the gradients of intermediates are not needed.
+    trace = thunder.core.transform_common.dce(trace)
+
+    end_time_ns = time.perf_counter_ns()
+    elapsed_time_ns = end_time_ns - start_time_ns
+    elapsed_time_millis = elapsed_time_ns // 1000000
+    trace.set_provenance(
+        thunder.core.trace.TraceProvenance(f"Grad transform pass (took {elapsed_time_millis} milliseconds)")
+    )
+    return trace
+
+
+def split_into_forward_and_backward(joint_trace):
+    """split a joint trace for forward and backward into separate ones, including recomputation (aka activation checkpointing)"""
+
+    # the joint trace will have the forward computation at the beginning and then the backward computation
+    # from how it is constructed.
+    # we split the trace:
+    # - forward symbols go into forward_part_bsyms
+    # - all symbols not in the forward go into backward_part_bsyms
+    # - for recomputation (aka activation checkpointing), we want to insert symbols going into the forward also into the
+    #   backward. but we want to do so "just in time". To this end, we gather the symbols in a dict and later
+    #   insert it when their respective outputs are needed. This is in backward_part_bsyms_recomputated
+    #   The just in time recomputation is a heuristic to save memory mimicking checkpointing: e.g. for a checkpointed
+    #   block, the forward would be recomputed just before computing the gradient.
+    # the splitting is done in the reverse order fo the bound symbols, and works out which bits are needed for the forward
+    # from there.
+    forward_part_bsyms = []
+    backward_part_bsyms = []
+    backward_part_bsyms_recomputed: dict[str, tuple[BoundSymbol, set[str]]] = {}  # name -> (bsym, name_set)
+
+    # return has a dict, it contains the forward outputs in "output" and the backward outputs in "grad_flat_args", corresponding
+    # to "flat_args"
+    return_bsym = joint_trace.bound_symbols[-1]
+    assert return_bsym.sym == prims.python_return
+    fw_output = return_bsym.args[0]["output"]
+    assert isinstance(fw_output, tuple)
+
+    grad_outs = [None for _ in fw_output]
+    output_pos = {o.name: i for i, o in enumerate(fw_output) if isinstance(o, thunder.TensorProxy)}
+
+    # the proxies we need to compute in the forward - we start with the outputs of the forward
+    forward_proxy_names = {o.name for o in thunder.core.pytree.tree_iter(fw_output) if isinstance(o, thunder.Proxy)}
+    # we also have the inputs available, so we add flat_args.
+    # for inplace, we need to update this (or have flat args be the right thing?...)
+    forward_proxy_names.update(a.name for a in return_bsym.args[0]["flat_args"] if isinstance(a, thunder.Proxy))
+
+    # We keep track of the names of proxies we recompute in the backward as those will not need to be part of the
+    # ones saved in the forward for the backward
+    backward_recomputed_proxy_names = set()
+
+    # here we go through the bound symbols to assign to the three categories, keeping track of the forward names, the backward needed
+    # and the backward available names.
+    for bsym in reversed(joint_trace.bound_symbols):
+        if bsym.sym == prims.python_return:  # we will re-do the return statements below, so we skip it here
+            continue
+
+        # unpack_trivial is always a (forward trace) argument, the backward inputs are from get_grad
+        # anything that outputs something needed for the forward (the forward_proxy_names) is part of the forward
+        if any((o.name in forward_proxy_names) for o in bsym.flat_proxy_outs) or bsym.sym == prims.unpack_trivial:
+            forward_part_bsyms.insert(0, bsym.from_bsym())
+            bsym_arg_names = {a.name for a in bsym.flat_proxy_args}
+            bsym_out_names = {o.name for o in bsym.flat_proxy_outs if o.name not in bsym_arg_names}
+            forward_proxy_names.update(bsym_arg_names)
+            forward_proxy_names.update(bsym_out_names)
+
+            # if we need to recompute this, we also add the bound symbol to backward_part_bsyms_recomputed
+            if _should_recompute_bsym_in_backward(bsym):
+                backward_recomputed_proxy_names.update(bsym_out_names)
+                bsym_rec = (bsym.from_bsym(), bsym_out_names)
+                backward_part_bsyms_recomputed.update((n, bsym_rec) for n in bsym_out_names)
+
+            continue
+
+        # get grad is always part of the input, record the grad_out (will be part of the "cotangents" list)
+        if bsym.sym == prims.get_grad:
+            grad_outs[output_pos[bsym.args[0].name]] = bsym.output
+            continue
+
+        # copy_ updating a forward proxy is special regardless of the output
+        if bsym.sym == prims.copy_ and bsym.args[1].name in forward_proxy_names:
+            # todo: should we also handle ltorch.copy_ ?
+            forward_part_bsyms.insert(0, bsym.from_bsym())
+            forward_proxy_names.update(a.name for a in bsym.flat_proxy_args)
+            continue
+
+        # if we don't need to have it in the forward, it is part of the backward
+        backward_part_bsyms.insert(0, bsym.from_bsym())
+
+    # we insert the recomputed symbols just before where they are needed.
+    # as this inserts into the list during processing, we use a while loop rather than
+    # a for loop
+    bw_idx = 0
+    while bw_idx < len(backward_part_bsyms):
+        bsym = backward_part_bsyms[bw_idx]
+        modified = False
+        for n in reversed(bsym.flat_proxy_args):
+            recomp_bsym_rec = backward_part_bsyms_recomputed.get(n.name)
+            if recomp_bsym_rec is not None:
+                modified = True
+                recomp_bsym, recomp_output = recomp_bsym_rec
+                backward_part_bsyms.insert(bw_idx, recomp_bsym)
+                for nn in recomp_output:
+                    del backward_part_bsyms_recomputed[nn]
+        if not modified:
+            bw_idx += 1
+
+    # collect proxies needing to be saved in the forward for the backward
+    saved_for_backward = {}
+    for bsym in backward_part_bsyms:
+        saved_for_backward.update(
+            (a.name, a)
+            for a in bsym.flat_proxy_args
+            if a.name in forward_proxy_names and a.name not in backward_recomputed_proxy_names
+        )
+    saved_for_backward_tensors = [p for p in saved_for_backward.values() if isinstance(p, thunder.TensorProxy)]
+    saved_for_backward_other = [p for p in saved_for_backward.values() if not isinstance(p, thunder.TensorProxy)]
+
+    # we build the forward trace
+    forward_trace = thunder.core.trace.from_trace(joint_trace)
+    forward_trace.tags.add(thunder.core.trace.TraceTag.AUGMENTED_FORWARD)
+    forward_trace.names = forward_trace.names.copy()  ## ehem
+    forward_trace.bound_symbols += forward_part_bsyms
+
+    # now we create the return value and return bound symbol for  the forward
path: thunder/transforms/autodiff.py
user: login: beverlylytle
body: ```suggestion
    # now we create the return value and return bound symbol for the forward
```
created_at: 2025-05-20T12:28:29Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,477 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims, utils
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+    recompute_saved_for_backward,
+)
+from thunder.core.proxies import ProxyTag
+from thunder.core.symbol import BoundSymbol
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+def _should_recompute_bsym_in_backward(bsym):
+    return any((ProxyTag.RECOMPUTE_IN_BACKWARD in o.tags) for o in bsym.flat_proxy_outs)
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                ########### BEGINNING of return handling (and putting the backward computation in the joint trace)
+                # This is big (and a bit messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols backward in order (from the generation of the list, this
+                # is the backward of the last forward instruction first, then then the backward of the second to last etc.
+                for bw_bsyms in self.collected_bw_part_bsyms:
+                    # if we don't have and grad_outs, we know that the grad_ins are not needed, so we don't compute them
+                    # (e.g. for `where(a * b > 0, a, b)` we do not need the grad components of a * b propagated
+                    if any(
+                        ((gg.name in grad_proxy_map) or (gg.name in output_proxy_names))
+                        for gg in (bs.args[0] for bs in bw_bsyms if bs.sym == prims.get_grad)
+                    ):
+                        for nbsym in bw_bsyms:
+                            if nbsym.sym == prims.put_grad:
+                                # put_grad adds something to the gradients
+                                # - if it is an output of the function, we also need to get the
+                                #   "grad_output" (but only once)
+                                # - if we already had a gradient, we need to add the new component
+                                # - structurally (for correctness), there should be no put_grad
path: thunder/transforms/autodiff.py
user: login: t-vi
body: yeah, let's do it.
created_at: 2025-05-20T12:51:05Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,477 @@
+import thunder.core.transforms
+from thunder.core.transforms import ForwardBackwardTraces
+
+from thunder.core import prims, utils
+from thunder.core.transforms import (
+    is_constant_for_vjp,
+    _get_gradfn_and_executor,
+    augmented_forward_impls,
+    backward_impls,
+    recompute_saved_for_backward,
+)
+from thunder.core.proxies import ProxyTag
+from thunder.core.symbol import BoundSymbol
+from thunder.core.vjp_utils import make_aug_forward_and_backward
+from thunder.core.pytree import tree_map
+import thunder
+import time
+
+
+def _should_recompute_bsym_in_backward(bsym):
+    return any((ProxyTag.RECOMPUTE_IN_BACKWARD in o.tags) for o in bsym.flat_proxy_outs)
+
+
+# Transforms a trace by determining which grad transforms to call given the list of executors in priority order
+# This pass tries to preserve the original trace and proxies.
+def grad_transform_on_trace(trace, /, *args, **kwargs):
+    # This processes the bsyms to map symbols to operator executors:
+    # - in the order of the executor list
+    #   - if the executor defines a grad transform, call that to
+    #     get fw + bw,
+    # - if there is a augmented forward registered, use that and the backwared registered,
+    #   to construct a grad transform to get fw + bw
+    # - if neither of the above apply, and the symbol has subsymbols, push the decomposition
+    #   to the front of the queue
+    # - if none of the above apply and we have a prim, raise an error
+    class AugmentedForwardProcessor(thunder.core.trace_interpreter.TraceSubstitutionProcessor):
+        def __init__(self, trace):
+            super().__init__(trace)
+            self.collected_bw_part_bsyms = []
+
+        def process_bsym(self, bsym: thunder.core.symbol.BoundSymbol) -> None:
+            if bsym.sym is prims.python_return:
+                ########### BEGINNING of return handling (and putting the backward computation in the joint trace)
+                # This is big (and a bit messy):
+                # the return (of the input trace) signals the end of the forward processing
+                # all the backwards will have been put in self.collected_bw_part_syms
+                # special symbols are
+                # - get_grad "reads" the gradient of a variable from the forward to compute with it
+                #   (if you are used to torch.autograd.Function, this is a "grad_out" in the arguments
+                #    of the backward static method)
+                # - put_grad "writes" a computed gradients (similar to returning the "grad_in" in t.a.F)
+                # The backward computation in self.collected_bw_part_syms is in the oder of th
+                # backward computation.
+                #
+                input_proxy_names = {p.name for p in bsym.args[0]["flat_args"] if isinstance(p, thunder.Proxy)}
+                output_proxy_names = set()
+                for o in thunder.core.pytree.tree_iter(bsym.args[0]["output"]):
+                    if isinstance(o, thunder.Proxy):
+                        output_proxy_names.add(self.read(o).name)
+                grad_proxy_map = {}
+
+                # we iterate through the collected symbols backward in order (from the generation of the list, this
+                # is the backward of the last forward instruction first, then then the backward of the second to last etc.
+                for bw_bsyms in self.collected_bw_part_bsyms:
+                    # if we don't have and grad_outs, we know that the grad_ins are not needed, so we don't compute them
+                    # (e.g. for `where(a * b > 0, a, b)` we do not need the grad components of a * b propagated
+                    if any(
+                        ((gg.name in grad_proxy_map) or (gg.name in output_proxy_names))
+                        for gg in (bs.args[0] for bs in bw_bsyms if bs.sym == prims.get_grad)
+                    ):
+                        for nbsym in bw_bsyms:
+                            if nbsym.sym == prims.put_grad:
+                                # put_grad adds something to the gradients
+                                # - if it is an output of the function, we also need to get the
+                                #   "grad_output" (but only once)
+                                # - if we already had a gradient, we need to add the new component
+                                # - structurally (for correctness), there should be no put_grad
+                                #   after a get_grad
+                                p = self.env.get(nbsym.args[0].name, nbsym.args[0])
+                                current_grad = grad_proxy_map.get(p.name)
+                                new_grad = nbsym.args[1]
+                                if current_grad is None and p.name in output_proxy_names:
+                                    self.new_trace.push_scope([])
path: thunder/transforms/autodiff.py
user: login: t-vi
body: Interesting. So 
```python
with self.new_trace.capture_in_scope() as new_bsyms:
   current_grad = prims.get_grad(p)
self.add_processed_bsyms(new_bsyms)
```
instead of
```python
self.new_trace.push_scope([])
current_grad = prims.get_grad(p)
self.add_processed_bsyms(self.new_trace.pop_scope())
```
?
It emphasizes the block nature.
created_at: 2025-05-20T12:55:11Z
author_association: COLLABORATOR

PR comments:
user: login: t-vi
created_at: 2025-04-13T15:47:57Z
body: Aside form not computing the gradient of process groups, there is an error in distributed stemming from an incomplete substitution of proxies in the subsymbols following synchronize. This is an example from an input trace in test_fsdp (`t_net2_weight` should have been replaced with `t22` also in the `prims.linear`):
```
  # /home/zeus/miniconda3/envs/cloudspace/lib/python3.12/site-packages/torch/nn/modules/linear.py:125:          return F.linear(input, self.weight, self.bias)
  t18 = ltorch.linear(input, t22, t21)  # t18: "cuda:1 f32[2, 2]"
    # t18 = prims.linear(input, t_net2_weight, t_net2_bias)  # t18: "cuda:1 f32[2, 2]"
```

I'll look at fixing the FSDP transform.

author_association: COLLABORATOR
user: login: t-vi
created_at: 2025-04-20T10:57:08Z
body: One failure is from the assumption that augmented forward and backward implementations would have flat inputs and outputs, an assumption that seems to be violated (only) by the automatically registered `multi_dot`.
I'd be inclined to not fix this but to either move auto registration to the grad_transform scheme or manually provide multi_dot.
I would likely be interesting to actually write a decomposition for multi_dot, as it reduces to a sequence of matmuls.
I opened #1975 for a reduction. 
author_association: COLLABORATOR
user: login: t-vi
created_at: 2025-05-04T20:40:40Z
body: CPU Tests are good now, as is GPU main, GPU ops and GPU grad.

GPU distributed
- [x] TensorParallel: 4x, one of them
  `FAILED thunder/tests/distributed/test_tensor_parallel.py::TensorParallelTest::test_both_column_and_row_bias_False`
  shape mismatch (maybe shapes not updated in transform #1179?)
- [x] `FAILED thunder/tests/distributed/test_fsdp.py::FSDPTest::test_rematerialize_all_gather`

author_association: COLLABORATOR
user: login: riccardofelluga
created_at: 2025-05-09T17:46:47Z
body: All pretty interesting stuff! As far as architecture goes, is the idea of trying to keep one single trace for as long as possible inside thunder and then split it only at the very end to pipe it in the torch autograd infrastructure? 

Maybe not in this PR but just asking to get this out to understand the direction. I think it would be pretty beneficial to do all the compiler passes with one trace, as it allows to reason with a wider scope(and later transformations can take advantage of this one single trace)
author_association: COLLABORATOR
user: login: t-vi
created_at: 2025-05-12T12:02:23Z
body: > is the idea of trying to keep one single trace for as long as possible inside thunder and then split it only at the very end to pipe it in the torch autograd infrastructure

So (quite) a few of us @beverlylytle @IvanYashchuk @lantiga chatted about this.

There are main two parts we want to evolve the autodiff:
- this PR produces a joint trace (and changes the way we produce it), but for the moment splits it very soon (as before),
- @beverlylytle is working on keeping a single trace until after transform for execution.
author_association: COLLABORATOR
user: login: t-vi
created_at: 2025-05-19T10:41:50Z
body: I'll split out some of the stand-alone features.
author_association: COLLABORATOR
user: login: t-vi
created_at: 2025-05-19T13:18:45Z
body: I am factoring out fixes, so I guess this will need to merge after them (and after merging main).
author_association: COLLABORATOR
user: login: t-vi
created_at: 2025-05-20T13:31:18Z
body: Thank you @beverlylytle @riccardofelluga @IvanYashchuk @lantiga 
author_association: COLLABORATOR