Hitchhiker's Guide — PyTorch & TensorFlow Framework Internals

"You don't understand a system until you've broken it." — every production on-call engineer

This guide goes beneath the public API. It is the answer to "why does this work?" when the official docs say "just call .backward()."

1. PyTorch Execution: The Full Stack

1.1 From Python to Kernel

When you write y = torch.matmul(a, b) in Python, here is the exact call chain:

Python: torch.matmul(a, b)
   → C++: at::matmul(self, other)                          # aten namespace
   → c10::Dispatcher::call(op_handle, stack)               # dispatch table lookup
   → Dispatch key selection: max(a.key_set(), b.key_set()) # e.g., CUDA
   → at::native::mm_cuda(self, other)                      # CUDA kernel dispatch
   → cublasSgemm / cublasGemmEx                            # vendor BLAS

The dispatcher (c10::Dispatcher) is PyTorch's universal function registry. Every operator is registered as an entry in the dispatcher table, keyed by a DispatchKey. The dispatcher picks the highest-priority key in the tensor's DispatchKeySet.

1.2 Dispatch Key Priority (highest to lowest)

Functionalize          ← wraps mutations as functional ops (for torch.compile)
InferenceMode          ← disables autograd
AutogradMeta           ← for meta/fake tensors in torch.compile
Autograd               ← wraps kernel in grad recording logic
AutocastCPU/CUDA       ← AMP dtype promotion
CUDA                   ← GPU kernels
CPU                    ← CPU kernels
Meta                   ← shape-only "meta" device (used in compilation)

Why this matters for NPU deployment: When you register a custom NPU dispatch key (e.g., Qualcomm's QnnBackend), it must be inserted at the right priority. Below Autograd means the kernel won't record gradients — which is correct for inference-only NPU deployments.

1.3 Autograd Engine Deep Dive

Forward pass — DAG construction:

a = torch.tensor([2.0], requires_grad=True)
b = torch.tensor([3.0], requires_grad=True)
c = a * b   # grad_fn = MulBackward0
d = c + a   # grad_fn = AddBackward0

After these four lines, the autograd engine has built:

d  (AddBackward0)
├── c  (MulBackward0)
│   ├── a  (leaf)
│   └── b  (leaf)
└── a  (leaf)

Each non-leaf tensor has a grad_fn that stores:

1. The saved tensors needed for the backward pass
2. References to the next_functions (upstream nodes in the DAG)

Backward pass — Engine::execute():

The backward pass is triggered by d.backward(). The engine:

1. Starts from d.grad_fn with initial gradient torch.ones(1)
2. Calls AccumulateGrad for leaf tensors
3. Calls each grad_fn.__call__(gradient_from_downstream) in reverse topological order
4. Accumulates gradients via add_ into each leaf's .grad

Why topological order? A node can only compute its upstream gradient after it receives all gradients from downstream nodes that depend on it. This is Dijkstra's algorithm in reverse — we process nodes when their contribution count reaches zero.

Critical invariant: grad_fn is attached during the forward pass. The backward graph is fully determined by what ops ran in forward. This is why torch.no_grad() is a massive speedup — no Function objects are allocated.

1.4 torch.autograd.Function — The Correct Mental Model

class ClampedSoftmax(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, min_val, max_val):
        output = torch.softmax(input, dim=-1).clamp(min_val, max_val)
        ctx.save_for_backward(output)           # saves tensor refs (not copies)
        ctx.min_val = min_val                   # non-tensor metadata OK
        return output

    @staticmethod
    def backward(ctx, grad_output):
        output, = ctx.saved_tensors
        # Jacobian of softmax: J_ij = o_i * (delta_ij - o_j)
        # With clamping: gradient is zero where output was clamped
        mask = (output > ctx.min_val) & (output < ctx.max_val)
        # Softmax backward: grad_input = output * (grad_output - (grad_output * output).sum(-1, keepdim=True))
        s = (grad_output * output).sum(-1, keepdim=True)
        grad_input = output * (grad_output - s)
        grad_input = grad_input * mask.float()
        return grad_input, None, None  # None for non-tensor args

Common pitfalls:

1. ctx.save_for_backward(x) → storing x directly in ctx (leaks memory, misses memory-efficient inplace handling)
2. Returning wrong number of gradients — must match number of forward inputs exactly (use None for non-differentiable)
3. create_graph=True — if your backward uses non-differentiable ops, higher-order gradients will fail silently with wrong values

2. torch.fx — Symbolic Tracing Mechanics

2.1 What FX Actually Does

torch.fx.symbolic_trace(model) runs the model's forward() with Proxy objects instead of real tensors. Every operation on a Proxy records a Node in a Graph:

import torch.fx as fx

class MyModel(torch.nn.Module):
    def forward(self, x):
        return torch.relu(x + 1)

model = MyModel()
graph_module = fx.symbolic_trace(model)
print(graph_module.graph)
# graph():
#   %x : [#users=1] = placeholder[target=x]
#   %add : [#users=1] = call_function[target=operator.add](args = (%x, 1))
#   %relu : [#users=1] = call_function[target=torch.relu](args = (%add,))
#   return relu

The result is an FX IR — a data structure you can inspect, modify, and re-compile.

2.2 FX Graph Transformation — The Correct Pattern

class FuseConvBN(fx.Transformer):
    def call_module(self, target, args, kwargs):
        # Check if this is a Conv2d followed by BatchNorm2d
        module = self.fetch_attr(target)
        if isinstance(module, torch.nn.Conv2d):
            # Look ahead in graph — check next user
            node = self.current_node
            if len(node.users) == 1:
                next_node = list(node.users)[0]
                if (next_node.op == 'call_module' and
                    isinstance(self.fetch_attr(next_node.target), torch.nn.BatchNorm2d)):
                    # Fuse: fold BN params into Conv weights
                    fused = torch.nn.utils.fusion.fuse_conv_bn_eval(
                        module, self.fetch_attr(next_node.target)
                    )
                    # Register fused module, return its output
                    self.root.add_module(target + '_fused', fused)
                    return self.tracer.call_module(target + '_fused', args, kwargs)
        return super().call_module(target, args, kwargs)

2.3 FX Limitations You MUST Know for Interviews

1. No data-dependent control flow: if x.sum() > 0: cannot be traced — the Proxy doesn't know the value
2. No dynamic shapes by default: shapes are captured at trace time and become constants
3. Closures and globals: Python closures are NOT captured in the graph — they become constants at trace time, which is a semantic footgun
4. torch.compile vs FX: torch.compile uses TorchDynamo (bytecode analysis) instead of symbolic tracing — it CAN handle control flow by creating multiple graphs with guards

3. torch.compile Stack

3.1 Three-Layer Architecture

User Python code
     ↓
TorchDynamo           ← bytecode analysis, guards, graph breaks
     ↓  (FX Graph)
AOTAutograd           ← joint forward+backward graph
     ↓  (FX Graph)
TorchInductor         ← code generation (Triton for GPU, C++ for CPU)
     ↓
Compiled artifact     ← .so or PTX or Metal shader

3.2 TorchDynamo — How Bytecode Analysis Works

TorchDynamo installs a frame evaluation callback that intercepts Python bytecode execution. When it encounters a function decorated with torch.compile, it:

1. Disassembles the bytecode with dis.get_instructions()
2. Tracks every operation that produces a tensor
3. When it hits an unsupported operation (Python control flow on tensor values), it breaks the graph and falls back to Python
4. For captured graph segments, it installs guards — runtime checks that verify the assumptions made during tracing are still valid

@torch.compile
def f(x):
    if x.shape[0] > 10:  # <- GUARD: shape[0] > 10
        return x * 2
    return x + 1

First call with x.shape = (15,) → compiles graph 1 (x * 2) with guard shape[0] > 10
Second call with x.shape = (5,) → guard fails → compiles graph 2 (x + 1) with guard shape[0] <= 10
Third call with x.shape = (20,) → guard 1 passes → reuses compiled graph 1

3.3 Custom Backend Registration

from torch._dynamo import register_backend

@register_backend
def my_npu_backend(gm: torch.fx.GraphModule, example_inputs):
    # gm is the FX graph module
    # example_inputs is a list of example tensors for shape/dtype inference
    
    # Step 1: Optimize FX graph
    gm = optimize_for_npu(gm)
    
    # Step 2: Lower to NPU IR
    npu_program = lower_to_npu_ir(gm, example_inputs)
    
    # Step 3: Return a callable
    def run(*args):
        return npu_program.execute(args)
    
    return run

model_compiled = torch.compile(model, backend="my_npu_backend")

4. TensorFlow Graph Execution Mechanics

4.1 tf.function Tracing Protocol

When Python hits a @tf.function-decorated function for the first time:

1. Tracing: Python executes with SymbolicTensor objects (like PyTorch Proxies)
2. AutoGraph transformation: Python if/for/while are converted to tf.cond/tf.while_loop
3. Graph construction: a FuncGraph is built representing the TF ops
4. Compilation: the graph is compiled to a ConcreteFunction with a specific input signature
5. Execution: subsequent calls with matching signatures execute the cached ConcreteFunction

4.2 Retracing Triggers — The Complete List

@tf.function
def f(x, training=False):
    return x * 2 if training else x + 1

f(tf.ones([3]))          # trace 1: input_signature=(TensorSpec(shape=[3], dtype=float32),), training=False
f(tf.ones([3]))          # reuse trace 1 ✓
f(tf.ones([4]))          # trace 2: new shape [4] ← RETRACE (no input_signature specified)
f(tf.ones([3]), training=True)  # trace 3: new Python value for training ← RETRACE
f(tf.constant(1.0))     # trace 4: scalar shape ← RETRACE

Prevention: use @tf.function(input_signature=[tf.TensorSpec(shape=[None, 128], dtype=tf.float32)]) to fix dtype/rank and allow variable batch/sequence dimensions.

4.3 tf.Variable Inside tf.function

# WRONG — creates a new Variable on every trace:
@tf.function
def bad(x):
    v = tf.Variable(0.0)  # ← Python-side effect, runs at trace time only!
    v.assign_add(x.sum())
    return v

# CORRECT — Variable must be created outside tf.function:
v = tf.Variable(0.0)
@tf.function
def good(x):
    v.assign_add(x.sum())
    return v

5. Numerical Precision Deep Dive

5.1 FP32 vs BF16 vs FP16

Key insight for NPU work: BF16 has the same dynamic range as FP32 (same exponent bits) but lower precision. This makes BF16 safer for training (no overflow/underflow) but potentially less accurate for narrow distribution outputs. Qualcomm's HTP natively supports INT8/INT16/FP16/BF16 — knowing which precision per-layer gives the best accuracy-performance tradeoff is core to this role.

5.2 Catastrophic Cancellation and Kahan Summation

Summing [1e8, 1.0, -1e8] in FP32:

• Naive: (1e8 + 1.0) = 1e8 (1.0 is below precision of 1e8 in FP32), then 1e8 - 1e8 = 0.0 ← WRONG
• Kahan: maintains a running compensation term, gets 1.0 ← CORRECT

def kahan_sum(values):
    total = 0.0
    compensation = 0.0
    for v in values:
        y = v - compensation
        t = total + y
        compensation = (t - total) - y
        total = t
    return total

Interview answer: "Why does your softmax accumulate error in FP16?" → explain that the normalization sum in softmax has large magnitude (sum of exps), and the numerator single-exp is small — catastrophic cancellation. Fix: compute in FP32, cast output to FP16.

6. Common Interview Pitfalls

7. Resources & Further Reading

• PyTorch Internals Blog — Edward Yang
• TorchDynamo — How it Works (PyTorch Dev Blog)
• Autograd: Automatic Differentiation in ML (JMLR 2017) — the paper behind modern autodiff engines
• TF tf.function Guide
• PyTorch FX API Docs
• AOTAutograd internals
• Paper to read: "JAX: Composable Transformations of Python+NumPy Programs" (Bradbury et al., 2018) — JAX solves the same problems as torch.compile but differently; understanding both is a superpower