Tensor Core Implementation

Beyond the API

Frameworks like PyTorch and TensorFlow provide powerful, high-level tools for building neural networks. But what's happening under the hood? A PyTorch tensor isn't magic; it's a clever data structure. At its core, it's a view into a contiguous, one-dimensional block of memory. All the complexity of multiple dimensions is handled by a small set of metadata that tells the computer how to interpret that single list of numbers.

To truly understand how deep learning frameworks operate, we're going to build our own Tensor class from scratch using Python and NumPy. This will reveal how multidimensional data is managed, manipulated, and prepared for the calculations that drive machine learning.

Memory Layout and Strides

Imagine a 2x3 tensor containing the numbers 1 through 6. While we think of it as a grid, a computer's memory is just a long, flat line. The data is stored contiguously, like this: [1, 2, 3, 4, 5, 6].

So how does the computer know where the rows and columns are? It uses metadata called strides. Strides are a tuple of integers where each integer specifies the number of steps to take in memory to move one index along a particular dimension.

For our 2x3 tensor:

To move to the next row (dimension 0), you need to jump over 3 elements in memory. The stride for dimension 0 is 3.

To move to the next column (dimension 1), you just move to the next element. The stride for dimension 1 is 1.

The strides are (3, 1).

Another piece of metadata is the offset, which tells us where our tensor's data begins within the larger memory block. For now, it will be 0, but it becomes important when we create "views" of tensors, like slices, that point to the same underlying data but start at a different position.

The Tensor Class

Let's translate these ideas into a Python class. Our Tensor object will wrap a NumPy array, which will manage the underlying block of memory. We'll also store the shape, strides, and an offset. The grad attribute is included now to hold gradients for when we implement automatic differentiation later.

import numpy as np

class Tensor:
    def __init__(self, data, requires_grad=False):
        if not isinstance(data, np.ndarray):
            data = np.array(data, dtype=np.float32)
        
        # Core data and gradient
        self.data = data
        self.grad = None
        self.requires_grad = requires_grad

        # Metadata for views
        self.shape = data.shape
        self.strides = data.strides
        self.offset = 0 # Offset into the data buffer

    def __repr__(self):
        return f"Tensor(shape={self.shape}, strides={self.strides})\n{self.data}"

Here, we use NumPy's shape and strides attributes directly. A float32 data type is a common choice in machine learning for balancing precision and memory usage. Now we can create a tensor and inspect its properties.

a = Tensor([[1, 2, 3], [4, 5, 6]])
print(a)

# Output:
# Tensor(shape=(2, 3), strides=(24, 8))
# [[1. 2. 3.]
#  [4. 5. 6.]]

Making Tensors Work Together

A tensor class isn't useful if you can't perform operations on it. We can make standard Python operators like + and * work with our class by implementing special methods like __add__ and __mul__. This is called and allows for intuitive, mathematical syntax.

For a simple element-wise addition, we just add the underlying NumPy arrays and return a new Tensor containing the result.

class Tensor:
    # ... (previous __init__ and __repr__) ...

    def __add__(self, other):
        # Ensure 'other' is also a Tensor for the operation
        if not isinstance(other, Tensor):
            other = Tensor(other)
            
        # Perform the operation on the data arrays
        result_data = self.data + other.data
        return Tensor(result_data)

    def __mul__(self, other):
        if not isinstance(other, Tensor):
            other = Tensor(other)
        
        result_data = self.data * other.data
        return Tensor(result_data)

With this, we can now write clean, simple code:

a = Tensor([1, 2, 3])
b = Tensor([4, 5, 6])
c = a + b # Calls a.__add__(b)

print(c.data) # Output: [5. 7. 9.]

This works because NumPy itself handles the element-wise operation, including a powerful feature called which allows operations on arrays with different but compatible shapes. For example, you can add a vector to each row of a matrix.

matrix = Tensor([[10, 20, 30], [40, 50, 60]])
vector = Tensor([1, 2, 3])

result = matrix + vector
print(result.data)

# Output:
# [[11. 22. 33.]
#  [41. 52. 63.]]

We've laid the essential groundwork for a custom tensor library. We have a data structure that understands its layout in memory and can perform basic element-wise arithmetic. This foundation is what all higher-level operations, and ultimately neural networks, are built upon.

PyTorch Internal Mechanics and Implementation