compile_to_callable resolves Derivative nodes automatically

Users no longer call simplify/simplify_for_compile before compiling:

    step = compile_to_callable(loss, *[loss.grad(p) for p in params])

CompiledProgram detects Derivative nodes (simplify="auto") and applies the
compile-simplify preset internally with ONE shared memo across all outputs —
strictly better than per-output simplify_for_compile calls, since the loss
and its gradients keep their common forward subtrees shared at the symbolic
level. Pre-simplified inputs are untouched (behavior identical); simplify=
True/False overrides. compile_simplify_args() exposes the shared preset.

Suite: 782 passed.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: thomasahle

tensorgrad/compiler/runtime.py

@@ -17,16 +17,52 @@
 import sympy
 import torch
 
-from tensorgrad.tensor import Tensor, Variable
+from tensorgrad.tensor import Derivative, Tensor, Variable, compile_simplify_args
 from tensorgrad.compiler.lower import lower_program
 from tensorgrad.compiler.codegen_torch import TorchCodegen
 
 # Max number of cached shape/dtype specializations per program (LRU-evicted).
 SPECIALIZATION_CACHE_SIZE = 32
 
 
+def _contains_derivative(tensors) -> bool:
+    """Iterative DAG walk (shared visited-set across all outputs)."""
+    seen: set[int] = set()
+    stack = list(tensors)
+    while stack:
+        t = stack.pop()
+        if id(t) in seen:
+            continue
+        seen.add(id(t))
+        if isinstance(t, Derivative):
+            return True
+        stack.extend(getattr(t, "factors", ()) or ())
+        stack.extend(getattr(t, "terms", ()) or ())
+        stack.extend(getattr(t, "inputs", ()) or ())
+        if (inner := getattr(t, "tensor", None)) is not None:
+            stack.append(inner)
+    return False
+
+
 class CompiledProgram:
-    def __init__(self, tensors: tuple[Tensor, ...], verbose: bool = False, torch_compile: bool = False):
+    def __init__(
+        self,
+        tensors: tuple[Tensor, ...],
+        verbose: bool = False,
+        torch_compile: bool = False,
+        simplify: bool | str = "auto",
+    ):
+        # Derivative nodes can't be lowered; resolve them here so users can
+        # write compile_to_callable(loss, *[loss.grad(p) for p in params])
+        # directly. One shared args dict = one shared memo across ALL outputs,
+        # which preserves cross-output subtree sharing (better than per-output
+        # simplify_for_compile calls). "auto" simplifies only when a Derivative
+        # is present, keeping behavior identical for pre-simplified inputs.
+        if simplify == "auto":
+            simplify = _contains_derivative(tensors)
+        if simplify:
+            shared_args = compile_simplify_args()
+            tensors = tuple(t.simplify_for_compile(shared_args) for t in tensors)
         self.tensors = tensors
         self.verbose = verbose
         self.torch_compile = torch_compile
@@ -132,10 +168,22 @@ def __call__(self, values: dict, shapes: dict = None):
         return tuple(wrapped)
 
 
-def compile_to_callable(*tensors: Tensor, verbose: bool = False, torch_compile: bool = False):
+def compile_to_callable(
+    *tensors: Tensor,
+    verbose: bool = False,
+    torch_compile: bool = False,
+    simplify: bool | str = "auto",
+):
     """Compile one or more tensorgrad tensors into a fast callable.
 
+    Derivative nodes are resolved automatically (see simplify_for_compile), so
+    gradients can be passed raw:
+
+        step = compile_to_callable(loss, *[loss.grad(p) for p in params])
+
     Returns f(values: dict[Variable, torch.Tensor], shapes: dict[Symbol, int])
     -> named torch.Tensor or tuple of them (one per input tensor).
     """
-    return CompiledProgram(tuple(tensors), verbose=verbose, torch_compile=torch_compile)
+    return CompiledProgram(
+        tuple(tensors), verbose=verbose, torch_compile=torch_compile, simplify=simplify
+    )

tensorgrad/tensor.py

@@ -14,6 +14,22 @@
 from tensorgrad.utils import _MatchEdgesKey
 
 
+def compile_simplify_args() -> dict[str, Any]:
+    """The simplify() argument preset used on the way into the compiler:
+    resolve Derivative nodes, leave the algebra to the compiler's IR passes,
+    memoize across shared subtrees. One dict instance = one shared memo, so
+    passing the same dict to several simplify() calls preserves cross-output
+    subtree sharing (used by compile_to_callable for loss + gradients)."""
+    return {
+        "grad_steps": float("inf"),
+        "expand_functions": False,
+        "combine_products": False,
+        "sum_combine_terms": False,
+        "factor_components": False,
+        "memoize": True,
+    }
+
+
 _lazy_rename = False
 
 # Lazily imported tensorgrad.compiler.canon module (compositional structural
@@ -252,7 +268,7 @@ def full_simplify(self, expand=True) -> "Tensor":
                 expr = new
         return expr
 
-    def simplify_for_compile(self) -> "Tensor":
+    def simplify_for_compile(self, args: Optional[dict[str, Any]] = None) -> "Tensor":
         """Eliminate Derivative nodes with minimal other rewriting, so the
         compiler's IR passes (factoring, stabilization) do the algebra instead.
 
@@ -264,17 +280,15 @@ def simplify_for_compile(self) -> "Tensor":
         recover identical compact code at compile time — validated to match the
         full path's output while turning a >300s deep-model simplify into
         sub-second. Memoized by structural key across shared subtrees.
+
+        Note: `compile_to_callable` applies this automatically to any input
+        containing Derivative nodes, sharing one memo across all outputs — so
+        `compile_to_callable(loss, *[loss.grad(p) for p in params])` just
+        works. Call this directly only when you want the simplified symbolic
+        form itself. Pass `args` (from `compile_simplify_args()`) to share a
+        memo across several calls.
         """
-        return self.simplify(
-            {
-                "grad_steps": float("inf"),
-                "expand_functions": False,
-                "combine_products": False,
-                "sum_combine_terms": False,
-                "factor_components": False,
-                "memoize": True,
-            }
-        )
+        return self.simplify(args if args is not None else compile_simplify_args())
 
     @final
     def substitute(self, x: "Variable", y: "Tensor") -> "Tensor":

tests/test_simplify_memo.py

@@ -144,3 +144,36 @@ def grad_nodes(depth):
     n2, n4, n8 = grad_nodes(2), grad_nodes(4), grad_nodes(8)
     # Grows linearly, not exponentially: doubling depth must not >4x the nodes.
     assert n4 < 4 * n2 and n8 < 4 * n4, f"superlinear: {n2}, {n4}, {n8}"
+
+
+def test_compile_to_callable_accepts_raw_derivatives():
+    """compile_to_callable resolves Derivative nodes internally (shared memo
+    across outputs), so raw loss.grad(p) can be passed directly."""
+    b, i, j, k = symbols("b i j k")
+    x = Variable("x", b, i)
+    W1 = Variable("W1", i, j)
+    W2 = Variable("W2", j, k)
+    y = Variable("y", b, k)
+    loss = F.sum((F.relu(x @ W1) @ W2 - y) ** 2)
+    params = [W1, W2]
+
+    # The killer-demo one-liner: raw Derivative nodes straight in.
+    f_raw = compile_to_callable(loss, *[loss.grad(p) for p in params])
+    # Explicit route for comparison.
+    f_exp = compile_to_callable(loss.simplify_for_compile(),
+                                *[loss.grad(p).simplify_for_compile() for p in params])
+
+    dims = {b: 4, i: 5, j: 6, k: 3}
+    torch.manual_seed(0)
+    vals = {
+        x: torch.randn(4, 5).refine_names("b", "i"),
+        W1: torch.randn(5, 6).refine_names("i", "j"),
+        W2: torch.randn(6, 3).refine_names("j", "k"),
+        y: torch.randn(4, 3).refine_names("b", "k"),
+    }
+    outs_raw = f_raw(dict(vals), dims)
+    outs_exp = f_exp(dict(vals), dims)
+    for a, e in zip(outs_raw, outs_exp):
+        torch.testing.assert_close(
+            a.rename(None), e.align_to(*a.names).rename(None), rtol=1e-4, atol=1e-6
+        )