rafan
← back to blog

What Happens Below a Machine Learning Framework?

Tensors, autodiff, and Compiler Representations on a High-level

SEP 4 · 2026 · 14 MIN READ

Image from ML Frameworks and Abstractions slideshow
Image from Tianqi Chen's slides on ML Frameworks and Abstractions from Carnegie Mellon School of Computer Science

Prerequisite knowledge is required to understand every part of this article in full. Fret not, as I will do my best to provide in-text links and citations for much of the information within this article, including the prerequisite knowledge, which some won’t formally be in the references section. The in-text citations may be broad or vague regarding the definition of the hyperlinked word, but are connected to the context clues within this article when researched. All notation shown is original unless cited otherwise.

If you have ever navigated and implemented the process of building, training, and subsequent deployment of a deep learning model or neural network, one of the most accessible and common ways is to utilize what are called general libraries, which are programming libraries that one can import into their codebase that provide pre-made and ready-for-use code.1 Examples of ML frameworks include PyTorch, TensorFlow, and JAX. These frameworks provide a common toolset for the user and programmer to express mathematical intention, since the AI tasks, algorithms, and models are pre-built, without having to manage the actual machinery needed to do the math on the hardware directly.2

To demonstrate these abstractions, we will construct a simple neural network’s forward propagation and backpropagation mechanisms:

If we want to take an input of data xx and a set of learned knobs and configurations, i.e., weights WW that control how the input xx is transformed, we can then yield the neural network’s prediction, yy.1 This is denoted as

y=x×Wy = x \times W

For matrix multiplication to occur, the dimensions of two matrices must align. From linear algebra fundamentals, this means the number of columns in the first matrix must match the number of rows in the second.

So, with the data xx and the weight WW possessing the attribute of being an m×km \times k-sized matrix and a k×nk \times n-sized matrix, respectively, this expression can be extended to

x∈Rm×k,W∈Rk×n,y∈Rm×nx \in \mathbb{R}^{m \times k}, \qquad W \in \mathbb{R}^{k \times n}, \qquad y \in \mathbb{R}^{m \times n}

Thus, if we fully expand the notation of these matrices out, the multiplication actually looks like

[y11y12⋯y1ny21y22⋯y2n⋮⋮⋱⋮ym1ym2⋯ymn]⏟y∈Rm×n=[x11x12⋯x1kx21x22⋯x2k⋮⋮⋱⋮xm1xm2⋯xmk]⏟x∈Rm×k×[W11W12⋯W1nW21W22⋯W2n⋮⋮⋱⋮Wk1Wk2⋯Wkn]⏟W∈Rk×n\underbrace{ \begin{bmatrix} y_{11} & y_{12} & \cdots & y_{1n} \\ y_{21} & y_{22} & \cdots & y_{2n} \\ \vdots & \vdots & \ddots & \vdots \\ y_{m1} & y_{m2} & \cdots & y_{mn} \end{bmatrix} }_{y \in \mathbb{R}^{m \times n}} = \underbrace{ \begin{bmatrix} x_{11} & x_{12} & \cdots & x_{1k} \\ x_{21} & x_{22} & \cdots & x_{2k} \\ \vdots & \vdots & \ddots & \vdots \\ x_{m1} & x_{m2} & \cdots & x_{mk} \end{bmatrix} }_{x \in \mathbb{R}^{m \times k}} \times \underbrace{ \begin{bmatrix} W_{11} & W_{12} & \cdots & W_{1n} \\ W_{21} & W_{22} & \cdots & W_{2n} \\ \vdots & \vdots & \ddots & \vdots \\ W_{k1} & W_{k2} & \cdots & W_{kn} \end{bmatrix} }_{W \in \mathbb{R}^{k \times n}}

More accurately, in ML frameworks and fundamentals, the arbitrarily sized data matrix xx, weight matrix WW, and output/prediction matrix yy are better understood as tensors (which serves as an important definition and semantic backbone later on, as you will see). The number of dimensions in a tensor is called the rank.

In our previous matrices x∈Rm×kx \in \mathbb{R}^{m \times k}, W∈Rk×nW \in \mathbb{R}^{k \times n}, and y∈Rm×ny \in \mathbb{R}^{m \times n}, since there were two dimensions — m×km \times k, k×nk \times n and m×nm \times n — per matrix, these matrices would be categorized as rank-2 tensors.

And, since tensors can have arbitrarily sized dimensions, we can use rr, ss, and tt to denote the individual tensor ranks for xx, WW, and yy, respectively, and dd, ww, and ee to denote some assigned value/size of each dimension, also respectively. Therefore,

x∈Rd1×d2×⋯×dr,W∈Rw1×w2×⋯×ws,y∈Re1×e2×⋯×etx \in \mathbb{R}^{d_1 \times d_2 \times \cdots \times d_r}, \qquad W \in \mathbb{R}^{w_1 \times w_2 \times \cdots \times w_s}, \qquad y \in \mathbb{R}^{e_1 \times e_2 \times \cdots \times e_t}

2

In Python-based ML frameworks, this is commonly written using the @ operator, which signifies matrix multiplication operations and tensor contractions/manipulations.

y = x @ W

This is our linear transformation used in this simplified forward propagation mechanism.3

Now we have our prediction value from forward propagation, yy, at our disposal. However, we want our neural network to be able to learn things, so ideally, we would want to compare the value of the prediction yy to the correct answer of whatever the neural network is trying to learn, which we will instantiate as ytruey_{\text{true}}. This very comparison is called a loss function, which can be denoted as

L=loss(y,ytrue)\mathcal{L} = \text{loss}(y,y_{\text{true}})

Fundamentally, you can think of backpropagation as moving backwards from a forward propagation. We are trying to figure out which learned knobs contributed to how off the prediction yy was. This means we want to figure out how wrong the neural network’s prediction was and subsequently which learned knobs caused the network to be wrong and by how much.

So, we want to know how sensitive the loss function result L\mathcal{L} is to changes in the weight matrix WW by taking the partial derivative of L\mathcal{L} with respect to WW:

∂L∂W\frac{\partial \mathcal{L}}{\partial W}

And, if W∈Rk×nW \in \mathbb{R}^{k \times n}, this can be visualized as a rank-2 tensor (matrix) collection of partial derivatives, denoted as what can be referred to as the gradient of the loss4:

One way is to denote a single entry of the loss L\mathcal{L}'s local sensitivity to one particular weight Wi,jW_{i,j}, inside the gradient, such that

∂L∂Wi,j∈R\frac{\partial \mathcal{L}}{\partial W_{i,j}} \in \mathbb{R}

where ii denotes the row number of the matrix and jj denotes the column number. This makes Wi,jW_{i,j} a scalar value; a rank-0 tensor.

To articulate multiple ∂L∂Wi,j\frac{\partial \mathcal{L}}{\partial W_{i,j}} entries within the gradient that follows an assumption of W∈Rk×nW \in \mathbb{R}^{k \times n}, then

∂L∂W∈Rk×n\frac{\partial \mathcal{L}}{\partial W} \in \mathbb{R}^{k \times n}

A complete gradient containing the derivatives with respect to (w.r.t.) every component of WW is formalized in notation in a way that distinguishes generalizability regardless of the rank of WW, given the loss L\mathcal{L} is scalar-valued:

∇WL\nabla_W \mathcal L

Thus, in this W∈Rk×nW \in \mathbb{R}^{k \times n} paradigm, there exists an equivalence, such that

∂L∂Wij≡(∇WL)ij∈R∧∂L∂W≡∇WL∈Rk×n\frac{\partial \mathcal{L}}{\partial W_{ij}} \equiv \left(\nabla_W \mathcal{L}\right)_{ij} \in \mathbb{R} \quad \wedge \quad \frac{\partial \mathcal{L}}{\partial W} \equiv \nabla_W \mathcal{L} \in \mathbb{R}^{k \times n}

Therefore,

∇WL=[∂L∂W11∂L∂W12⋯∂L∂W1n∂L∂W21∂L∂W22⋯∂L∂W2n⋮⋮⋱⋮∂L∂Wk1∂L∂Wk2⋯∂L∂Wkn]\nabla_W \mathcal L = \begin{bmatrix} \frac{\partial \mathcal{L}}{\partial W_{11}} & \frac{\partial \mathcal{L}}{\partial W_{12}} & \cdots & \frac{\partial \mathcal{L}}{\partial W_{1n}} \\ \frac{\partial \mathcal{L}}{\partial W_{21}} & \frac{\partial \mathcal{L}}{\partial W_{22}} & \cdots & \frac{\partial \mathcal{L}}{\partial W_{2n}} \\ \vdots & \vdots & \ddots & \vdots \\ \frac{\partial \mathcal{L}}{\partial W_{k1}} & \frac{\partial \mathcal{L}}{\partial W_{k2}} & \cdots & \frac{\partial \mathcal{L}}{\partial W_{kn}} \end{bmatrix}

Remember that the loss L\mathcal{L} is not only yielded from changes in the weight matrix WW. In order for the prediction yy, it requires WW (forward propagation), and the prediction yy is important for computing loss(y,ytrue)\text{loss}(y, y_{\text{true}}). If you want to understand how WW influences L\mathcal{L}, the dependency path backwards could be thought of as

W→y→LW \rightarrow y \rightarrow \mathcal{L}

Following the line of dependencies — L\mathcal{L} depends on yy which depends on WW — we can invoke the chain rule from calculus, such that

∂L∂Wij=∑a=1m∑b=1n∂L∂yab∂yab∂Wij\frac{\partial \mathcal{L}}{\partial W_{ij}} = \sum_{a=1}^{m} \sum_{b=1}^{n} \frac{\partial \mathcal{L}}{\partial y_{ab}} \frac{\partial y_{ab}}{\partial W_{ij}}

Where aa indexes the mm rows of yy, and bb indexes the nn columns of yy.

To put it all together, if we take our forward propagation mechanism, y=x×Wy = x \times W, and we want to back-propagate/find the gradient of the loss w.r.t. WW, we will take the transpose of xx so that we maintain the matrix dimensions and align it for matrix multiplication, such that

xT∈Rk×mx^T \in \mathbb{R}^{k \times m}

So when

∇yL∈Rm×n\nabla_y \mathcal{L} \in \mathbb{R}^{m \times n}

then

∇WL=xT∇yL\nabla_W \mathcal{L} = x^T \nabla_y \mathcal{L}

And if we wanted to find the gradient of the loss w.r.t. data input xx, then it works the same way:

∇xL=∇yLWT\nabla_x \mathcal{L} = \nabla_y \mathcal{L} W^T

In PyTorch, this entire backward mechanism is triggered using a very simple line of code:

loss.backward()

In TensorFlow,

tf.GradientTape

And in JAX,

jax.grad

This is our backpropagation mechanism.

Already, off the bat, we can see how y = x @ W and loss.backward() are simple lines of code that act, in and of themselves, as an abstraction of the mathematics in formal articulation. For these lines of code, sprawling software-to-hardware relationships with foundational robustness must exist; y = x @ W can perform a forward pass with ease, and loss.backward() can elegantly trigger a series of backward-chained derivatives in a neural network. This is where the even deeper abstractions reside.

Underneath y = x @ W

PyTorch stores information like the data type, device, and layout when defining a tensor.3 Intuitively, the tensor’s shape tells the framework the dimensions of the data, and the layout and related metadata signal how that tensor is arranged and accessed in memory.

This is important when we analyze the perspective in which a piece of hardware takes upon a tensor’s presence: a x∈Rm×kx \in \mathbb{R}^{m \times k} matrix and a W∈Rk×nW \in \mathbb{R}^{k \times n} matrix have their values stored ultimately within a hardware system’s memory locations. This means that before an operation like @ can be validated and executed, a computer system must know whether the tensors are on a CPU, on an accelerator card, or in some other memory space. Within these constraints, the system must know a tensor’s dimensionality compatibility, memory layout, and whether it is suitable for the operation requested.

So, within the software-to-hardware dynamic, when y = x @ W is invoked, a framework must resolve what kind of tensor operation is being requested and choose an implementation based on it; a process that can be broadly understood as an operation dispatch. Depending on different tensor attributes, this operation dispatch can look different across the board.

We can get even lower than this; this was merely the software side of the hardware boundary here. There exists even more concreteness within the y = x @ W expression. As mentioned earlier, the matrix multiplication operation needs memory allocation, layout conversion, scheduling the processes and tasks around the operation, and hardware-specific kernels. As mentioned earlier, a great example is Tenstorrent’s open software stack, which separates the higher-level compiler and operation interfaces from the low-level programming and metal through the TT-Metalium repository.2, 4

As a wrap-up, not every ML framework adheres to the same philosophy, you could say, on how it handles something like Python code all the way down to hardware execution. For a system to turn y = x @ W into a concrete plan, it must understand the tensors and validate the operation invoked upon them. Then it chooses an implementation, manages the memory space, and then it finally performs the real work on the metal.

Underneath loss.backward()

The result of y = x @ W is a prediction value yy. Within a neural network, and its structure as previously stated, a comparison between the prediction and the correct answer ytruey_{\text{true}} is made via the loss L\mathcal{L}, yielded from the loss function loss(y,ytrue)\text{loss}(y,y_{\text{true}}). The process of backpropagation is to determine how changes in the learned weight values WW affect the loss, which is crucial for how a neural network actually learns; it tries to understand how much the internal learned knobs of the model contributed to the error of the output, and compute the gradients needed for the optimization step needed to update the weights WW.

Recall that when loss.backward() is invoked, the abstraction is actually the backwards computation of partial derivatives through the chain rule from calculus. ML frameworks must know which intermediate values the operations depend on, and how and in what order gradient information should be organized and subsequently propagated backwards through the network. This is where the concept of automatic differentiation, or autodiff, comes from.

Take micrograd as an example. Micrograd is a small, scalar-valued, reverse automatic differentiation engine (autograd) developed by Andrej Karpathy as an educational implementation of what lies underneath an ML framework during backpropagation. A dynamically constructed directed acyclic graph (DAG) is instantiated, which helps the engine remember the operations that produced the values from forward propagation, which previous values existed/depended on, and what local backward rule should be applied. In simpler terms, this all leaves a trail/record of how the loss value came to be.

This record-keeping structure can, at times, demonstrate that one value within this chain can affect the final loss value through multiple computational paths. The same weight WW can be used in multiple later operations, which means several downstream pieces can be affected when it comes to contributing toward that said weight’s final gradient. Before gradients are propagated backwards, the system has to traverse the DAG in an order that respects how each value was used. In Karpathy’s micrograd, it builds this topological ordering of a DAG before it applies the backward chain rules.

If we scale this ML framework infrastructure — to the size of PyTorch — the autograd mechanics follow the same broad idea. As tensor operations execute, PyTorch will construct this graph of backward Function objects. These intermediate values are saved and called upon when needed for future derivative calculations. Then, upon calling loss.backward(), autograd uses these operation history records to compute gradients. Ultimately, the goal is to preserve enough information during the forward pass so that the backwards pass can later apply the chain rule correctly.

In this sense, you could say that loss.backward() sits just one layer above the execution machinery that was discussed earlier. If you think of autograd as deciding what gradient work must happen, and the ML framework deciding how that work will run, on an accelerator-oriented stack, one could imagine it also eventually involving intimate tinkering with runtime execution, memory allocation, device-specific kernels, and the movement of tensor data through memory, just like forward pass operations require as well! 4, 2

Compiler Representations - High Level

Fundamentally, operation dispatch and execution models are related, but are separate decisions; operation dispatch selects an implementation for an operation, while execution models determine whether operations are issued immediately or captured into a larger representation.

Frameworks like PyTorch and TensorFlow, along with their compiler/runtime stacks, vary widely in how they handle the same line of y = x @ W code due to operation dispatch/implementation choices. Broadly, this means that the way PyTorch, TensorFlow, JAX, tinygrad, some TVM or MLIR-based systems, or hardware-specific stacks like Tenstorrent’s systems handle code, regarding operations and their subsequent implementations, is, unsurprisingly, different from one another.

For example, two different handling styles — eager execution systems and graph/compiler-oriented systems — represent two different execution styles:

  • In an eager execution system, an operation is issued when the program reaches it. This style is useful, intuitive, makes the program feel more direct, and can be debugged relatively easily. Simply: write an operation, run the operation, receive the result.

    • PyTorch is an example of a framework that uses this style by default; however, it also has graph/compiler-oriented systems as an option.

  • In graph/compiler-oriented systems, an operation may be captured into a larger representation before it is executed. This is because some systems retain a larger region of computation before an operation executes. This allows scheduling and lowering high-level tensor operations into lower-level representations and target-specific programs.

    • tinygrad is an example of a framework: they show documentation and an example of a system that moves from an abstracted Tensor API/frontend into UOps, scheduling, lowering, compilation, and runtime execution.

    • As an extension, there are broader compiler examples as well regarding retaining computation in intermediate representations (IR), which makes operations easier to analyze, transform, and subsequently optimize before it reaches the hardware.5, 6

In the exploratory pursuit of software-to-hardware dynamics, there exists a distinction between the gradient computation graphs and compiler IRs. You could think of an autograd graph as asking “how was this loss L\mathcal{L} produced, and how should the gradient ∇WL\nabla_W \mathcal L be computed?” The compiler IR would ask, “Well, how would this computation be represented, conditioned to the fact that it needs to be optimized and subsequently mapped into a target execution system?”

These roles are distinct, but they are not mutually exclusive: some systems perform automatic differentiation as a transformation over an intermediate representation, and computations produced by autodiff may later be captured and compiled.

As we descend deeper into the abstractions, and into the compiler representation side of things — i.e., intermediate representations — it is important to note that reverse-mode autodiff does not require the representation known as Static Single Assignment (SSA). In many compiler tasks, SSA can be useful; however, it is not a prerequisite for autodiff specifically. Micrograd is an example of this, since its reverse-mode autodiff is actually implemented via ordinary Python objects connected through the dynamic DAG rather than this compiler IR approach.7

Conclusion

To keep this article epistemically grounded, avoid scope creep, and not get too ahead of ourselves with the series of future articles succeeding this one — foundations for a much larger project — I want to conclude this piece. Already, our rudimentary instantiation of a basic forward and backward pass reveals the hidden layers of tensors, operation dispatch, compiler stuff, and autograd. Not every ML framework implements forward propagation or backpropagation the same way, with the same internal machinery.

References

Main References

  1. D. Zax, “Top Machine Learning Libraries,” IBM Think. [Online]. Available: https://www.ibm.com/think/topics/machine-learning-libraries. Accessed: Sep. 2, 2026.
  2. Tenstorrent, “Software Overview,” Tenstorrent Documentation. [Online]. Available: https://docs.tenstorrent.com/software/index.html. Accessed: Sep. 2, 2026.
  3. PyTorch, “Tensor Attributes,” PyTorch Documentation, ver. 2.13. [Online]. Available: https://docs.pytorch.org/docs/2.13/tensor_attributes.html. Accessed: Sep. 2, 2026.
  4. Tenstorrent, “TT-Metalium Documentation,” Tenstorrent Documentation. [Online]. Available: https://docs.tenstorrent.com/tt-metal/latest/tt-metalium/index.html. Accessed: Sep. 2, 2026.
  5. C. Lattner, M. Amini, U. Bondhugula, A. Cohen, A. Davis, J. Pienaar, R. Riddle, T. Shpeisman, N. Vasilache, and O. Zinenko, “MLIR: Scaling compiler infrastructure for domain specific computation,” in 2021 IEEE/ACM International Symposium on Code Generation and Optimization (CGO), 2021, pp. 2–14, doi: 10.1109/CGO51591.2021.9370308.
  6. T. Chen, T. Moreau, Z. Jiang, L. Zheng, E. Yan, H. Shen, M. Cowan, L. Wang, Y. Hu, L. Ceze, C. Guestrin, and A. Krishnamurthy, “TVM: An automated end-to-end optimizing compiler for deep learning,” in 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18), Carlsbad, CA, USA, 2018, pp. 578–594. [Online]. Available: https://www.usenix.org/system/files/osdi18-chen.pdf.
  7. A. Karpathy, “micrograd: A tiny scalar-valued autograd engine and a neural net library on top of it with PyTorch-like API,” GitHub. [Online]. Available: https://github.com/karpathy/micrograd. Accessed: Sep. 2, 2026.

Additional Sources

In-Text Links

Footnotes