How TPX separates computation and gradients

First answer: why place autograd in TPX?

P10 owns tensor values and operators. TPX owns the relationships between those values and the rules needed to compute gradients. If requires_grad, grad, and backward nodes lived directly in P10, every numerical program would carry training state and the compute and hardware layers would need to understand differentiation.

TPX instead composes p10::Tensor. Use tpx::Tensor when gradients are needed and call P10 directly for pure numerical work. The boundary lets a reader study how values are computed separately from how gradients are derived.

• P10: create tensors, execute operators, and select backend kernels.
• TPX: hold differentiation state, build the DAG, and schedule GradFn nodes.
• The boundary: connect the layers through tensor values and operator calls without merging their implementations.

On-demand tracing: build only the graph you need

When requires_grad is true, TPX records inputs, outputs, GradFn, and operator parameters, connecting each result to the node that produced it. When it is false, TPX does not create this backward state, so the forward path stays close to pure P10 execution.

This is more than a boolean switch. It determines which tensors retain context, which nodes enter the DAG, and which dependency path backward() can later traverse.

x = tp.ones((2, 2), requires_grad=True)
y = x.matmul(w).relu()

TPX records: matmul -> relu -> y
P10 executes: tensor kernels for matmul and relu

Separate forward computation from backward scheduling

During the forward pass, TPX only needs to record operator relationships. When backward() is called, it starts from the target node, topologically sorts the graph, and schedules each GradFn in reverse dependency order.

GradFn defines the local derivative rule, while P10 still executes the tensor additions, multiplications, and accumulations. Autograd rules and backend kernels can therefore evolve independently instead of being duplicated.

loss.backward()
  -> topologically sort the recorded DAG
  -> call each node's GradFn
  -> accumulate gradients with P10 operations
  -> expose x.grad and parameter gradients

Check a gradient through one concrete path

For y = x * w, the forward node keeps x, w, and y. During backward, the multiplication GradFn uses the upstream gradient to compute dy/dx = w and dy/dw = x, then relies on P10 for the tensor operations that accumulate those values.

Because the local rule is an explicit node, a numerical error can be narrowed to one GradFn, one accumulation, or one backend kernel instead of inferred only from the final loss.

• Inspect the inputs, outputs, and operator parameters held by each node.
• Check that the backward topological order matches the dependency graph.
• Use finite differences or a small hand calculation to validate local gradients.
• Compare forward values with tracing enabled and disabled.

Extension boundary: new rules, hardware, and experiments

A custom GradFn does not require changes to the P10 tensor core, and a new hardware backend does not need to reimplement the entire differentiation system. A researcher can replace one local rule and observe its effect on forward values, backward values, and memory independently.

That is the engineering value of TPX: it does not promise to solve every training-performance problem automatically. It makes differentiation an object that can be read, verified, and recomposed.