Tensor is the might

Every good abstraction solves a problem, and this post will cover everything I know so far about a brilliant math abstraction – tensors.

Neural networks, from a simple 2-layer MLP to GPT-5, all boil down to the same thing: floating-point numbers flowing through a graph of operations. This post builds a complete, accelerated tensor library from scratch in C. It is heavily inspired by Bellard’s libnc, which unfortunately has not been open-sourced yet.

A tensor is nothing but a flat array of numbers, plus some metadata telling you how to interpret those numbers as a multi-dimensional object. We all learned that 2D arrays can be better represented as 1D array plus a number of rows/columns – this is essentially what a tensor is.

But going beyond two dimensions – we might need some other metadata, such as a generalised shape:

float data[32 * 3 * 28 * 28]; // 32 images, 3 channels, 28x28 pixels
int shape[4] = ⚡; // shape of the tensor
int ndim = 4; // number of dimensions

Having a shape we can figure out that an element at position [n,c,h,w] in a 4D tensor lives at offset data[n*(3*28*28)+c*(28*28)+h*28+w], if we keep our tensor in a row-major C-order format (often the default in most tensor libraries today).

Now, calculating each time an index of an element like this is inefficient, so we can precompute the strides once the shape is known. Strides tell how many elements to skip if we want to advance for one element in the given dimension. We could also group all shape-related fields together:

struct ut_shape 💬;

struct ut_tensor 💬;

We can add some helpers to create a shape and get flat index from a multi-dimensional index:

ut_shape ut_shape_new(int ndims, int* dims) {
  ut_shape s = ⚡;
  for (int i = 0; i < ndim; i++) {
    s.shape[i] = dims[i];
    s.nelem *= dims[i];
  }
  return s;
}

int ut_index(ut_shape s, const int* idx) {
  int flat = 0, stride = 1;
  for (int i = s.ndim - 1; i >= 0; i--) {
    flat += idx[i] * stride;
    stride *= s.shape[i];
  }
  return flat;
}

ut_shape s = ut_shape_new(3, (int[]){2, 3, 4});
assert(s.nelem == 24);
assert(s.ndim  == 3);
// element [1][2][3] should be at offset 1*12 + 2*4 + 3 = 23
assert(ut_index(s, (int[]){1, 2, 3}) == 23);

Tensors are usually dynamically allocated, so we should provide a way to create and destroy them, nothing but a wrapper on top of malloc/free.

Sometimes we want to create a tensor that shares the same data with another tensor, for example to get a “view” of a part of a tensor, to transpose a tensor without copying data, or to modify its shape (flatten). In this case we need to track ownership of the data. It’d be also useful to “retain” tensors, so that they could outlive their original scope, like during iterative training to avoid the same allocation over and over. So, we add a reference counter field to the tensor struct:

struct ut_tensor {
  struct ut_shape shape; // shape of the tensor
  float *data; // pointer to the data
  int refcount; // reference count for shared ownership, free when reaches 0
  struct ut_tensor *owner; // pointer to the owner tensor if this tensor is a view
};

We can optimise memory management further, adding arena allocator or memory pool to avoid frequent malloc/free calls, but what we already have is a good start.

However, our tensors are hardly useful without operations on them.

Elementwise

The most basic operations on tensors are elementwise – a single loop over all elements, applying a function to each element (unary) or a pair of elements (binary). We can implement them just like that:

static void ew_neg(float* out, const float* a, int n) {
  for (int i = 0; i < n; i++) out[i] = -a[i];
}

// ...more unary ops...

static void ew_relu(float* out, const float* a, int n) {
  for (int i = 0; i < n; i++) out[i] = fmaxf(0.f, a[i]);
}
static void ew_add(float* out, const float* a, const float* b, int n) {
  for (int i = 0; i < n; i++) out[i] = a[i] + b[i];
}

// ...more binary ops...

static ut_tensor* ew_unary(ut_tensor* a, void (*fn)(float*, const float*, int)) {
  ut_tensor* out = ut_alloc(a->shape.ndim, a->shape.shape);
  fn(out->data, a->data, a->shape.nelem);
  return out;
}
static ut_tensor* ew_binary(ut_tensor* a, ut_tensor* b,
                            void (*fn)(float*, const float*, const float*, int)) {
  ut_tensor* out = ut_alloc(a->shape.ndim, a->shape.shape);
  fn(out->data, a->data, b->data, a->shape.nelem);
  return out;
}

ut_tensor* ut_neg(ut_tensor* a) { return ew_unary(a, ew_neg); }
ut_tensor* ut_exp(ut_tensor* a) { return ew_unary(a, ew_exp); }
ut_tensor* ut_sigmoid(ut_tensor* a) { return ew_unary(a, ew_sigmoid); }
ut_tensor* ut_tanh(ut_tensor* a) { return ew_unary(a, ew_tanh); }
ut_tensor* ut_relu(ut_tensor* a) { return ew_unary(a, ew_relu); }
ut_tensor* ut_add(ut_tensor* a, ut_tensor* b) { return ew_binary(a, b, ew_add); }
ut_tensor* ut_sub(ut_tensor* a, ut_tensor* b) { return ew_binary(a, b, ew_sub); }
ut_tensor* ut_mul(ut_tensor* a, ut_tensor* b) { return ew_binary(a, b, ew_mul); }
ut_tensor* ut_div(ut_tensor* a, ut_tensor* b) { return ew_binary(a, b, ew_div); }
ut_tensor* ut_scale(ut_tensor* a, float s) {
  ut_tensor* out = ut_alloc(a->shape.ndim, a->shape.shape);
  for (int i = 0; i < a->shape.nelem; i++) out->data[i] = a->data[i] * s;
  return out;
}

It’d be nice to assert() that the shapes of two tensors are the same before doing a binary operation. Alternatively, at this point we might decide that we want to support “broadcasting” – extending the smaller tensor to match the shape of the larger tensor. For example a tensor of shape [2, 3, 4] and a tensor of shape [3, 4] can be added together by “stretching” the second tensor along the first dimension. While being a common feature in many tensor libraries, I decided to leave it out for now. Most models I’m aiming for have their underlying tensors perfectly aligned. If not – we can either duplicate data explicitly aligning tensor shapes, or adding broadcasting index calculation later.

It might be tempting at this point to implement more operations, matrix multiplication, convolutions, etc. But this is the moment where we have to ask – do we want all operations to run on CPU?

Golden Processing Unit (GPU)

Looking at modern GPU prices, it is clear that they probably have a significant value when it comes to crunching arrays of numbers. Thus, instead of limiting our library to a CPU, we might consider offloading some operations to a GPU.

This is where things get interesting, because we now need to manage memory across both CPU and GPU, transfer data efficiently, learn how to write GPU kernels, and much more. To make things worse, there is no single “GPU accelerator” – there are many vendors, each with their own APIs and quirks: CUDA, OpenCL, WebGPU, Metal, Vulkan, etc.

I tried many times to get the most out of CPU-only tensors, but even with BLAS, LAPACK, OpenMP – I could not reach the same magnitude of performance as GPU-accelerated frameworks.

To keep things manageable, I decided to start with Metal. On the one hand it limits us to Apple devices, but on the other hand it’s a fairly modern API and simplifies things due to “unified memory” of Apple Silicon.

Metal uses its own language to define GPU kernels (MSL), which is similar to modern C++, and it compiles kernels at runtime, so we can define them as strings in our C code:

static const char *shaderSource =
    "#include \n"
    "using namespace metal;\n"
    "\n"
    "kernel void relu(device const float *in  [[buffer(0)]],\n"
    "                  device float       *out [[buffer(1)]],\n"
    "                  uint id [[thread_position_in_grid]]) {\n"
    "    float val = in[id];\n"
    "    out[id] = fmax(val, 0.0f);\n"
    "}\n";

To make this kernel run on GPU we need to create a device, build a command queue (GPUs are asynchronous), compile the kernel, create buffers for input and output, and finally dispatch the kernel to run on GPU. Here’s how a simple GPU-accelerated ReLU operation might look like:

// prepare device, queue, and compile kernel
id<MTLDevice> device = MTLCreateSystemDefaultDevice();
id<MTLCommandQueue> queue = [device newCommandQueue];
NSError *err = nil;
id<MTLLibrary> library = [device newLibraryWithSource:[NSString stringWithUTF8String:shaderSource] options:nil error:&err];
id<MTLFunction> reluFn = [library newFunctionWithName:@"relu"];
id<MTLComputePipelineState> pipeline = [device newComputePipelineStateWithFunction:reluFn error:&err];

// prepare data, dispatch kernel
float input[8]  = {-2.0, -1.0, 0.0, 1.0, 2.0, -0.5, 3.0, -3.0};
float output[8] = {0};
id<MTLBuffer> bufIn = [device newBufferWithBytes:input length:sizeof(input) options:MTLResourceStorageModeShared];
id<MTLBuffer> bufOut = [device newBufferWithLength:sizeof(output) options:MTLResourceStorageModeShared];
id<MTLCommandBuffer> cmdBuf = [queue commandBuffer];
id<MTLComputeCommandEncoder> enc = [cmdBuf computeCommandEncoder];
[enc setComputePipelineState:pipeline];
[enc setBuffer:bufIn  offset:0 atIndex:0];
[enc setBuffer:bufOut offset:0 atIndex:1];
MTLSize gridSize = MTLSizeMake(8, 1, 1);
MTLSize tgSize  = MTLSizeMake(pipeline.maxTotalThreadsPerThreadgroup, 1, 1);
[enc dispatchThreads:gridSize threadsPerThreadgroup:tgSize];
[enc endEncoding];
[cmdBuf commit];
[cmdBuf waitUntilCompleted];
// copy results back
memcpy(output, bufOut.contents, sizeof(output));

We can wrap the first part of this code into a “Metal context” struct, that is created once and is kept alive for the lifetime of the whole library. Then we can add helpers to allocate buffers, read and write them and dispatch kernels. While the official Metal API is in Objective-C or Swift, we can use low-level ObjC runtime functions to call it from C, to keep the library pure (objC_msgSend all over the code).

Now every operation on tensors has to be implemented twice – as a naïve CPU implementation, and as a GPU kernel. Depending on the origin of the tensor we would call one of the implementations. Note that this also means that tensors carry the “device” field, and we must specify the device when the tensor is allocated – CPU tensors would only have a normal data buffer, but Metal tensors would also have a pointer to a Metal buffer object.

It’s also becoming important that we need to track when GPU and CPU buffers diverge. An easy way is to keep “dirty” flags for CPU/GPU buffers and provide “sync” functions to copy data from one buffer to another.

Even though some operations can or should be done on a GPU, in some cases to simplify things we still create it on a CPU first, then do the processing and send/copy data back to GPU.

All in all, we’d need to implement the following:

static const char *_mtl_src = "..."; // MSL source code for all kernels
void *ut_mtl_init(); // create Metal context as a singleton
void ut_mtl_dispatch(void *ctx, char *kernel, void **bufs, int nbufs, void *bytes, int blen, int n); // dispatch kernel with buffers
void *ut_mtl_buf_alloc(void *ctx, void *data, int len); // allocate Metal buffer
void ut_mtl_buf_free(void *ctx, void *buf); // free Metal buffer
void ut_mtl_buf_read(void *ctx, void *buf, void *data, int len); // read Metal buffer to CPU
void ut_mtl_buf_write(void *ctx, void *buf, void *data, int len); // write CPU buffer to Metal

void ut_sync_cpu(ut_tensor *t); // sync CPU buffer from GPU
void ut_sync_gpu(ut_tensor *t); // sync GPU buffer from CPU
void ut_to_device(ut_tensor *t, ut_dev device); // move tensor to specified device (CPU/GPU)

Done with element-wise ops and their GPU kernel equivalents we can move on to the core of all neural networks – matrix multiplication.

Matmul

A made-up on the spot statistics tell me that 80% of FLOPs in any neural network are matrix multiplications. Getting them right and fast is crucial for any deep learning framework. There are many ways to implement matrix multiplication, the most efficient one depends on the size of the matrices, the hardware, and the memory layout.

First, let’s consider the simplest case of multiplying two 2D matrices A and B, where A is of shape [M, K] and B is of shape [K, N]. The result C will be of shape [M, N]. The naïve implementation could look like this:

void gemm(const float *A, const float *B, float *C,
          int m, int n, int k, bool ta, bool tb) {
    for (int i = 0; i < m; i++)
        for (int j = 0; j < n; j++) {
            float sum = 0.f;
            for (int l = 0; l < k; l++) {
                float a = ta ? A[l*m + i] : A[i*k + l];
                float b = tb ? B[j*k + l] : B[l*n + j];
                sum += a * b;
            }
            C[i*n + j] = sum;
        }
}

Needless to say this would be terribly slow – three nested loops and no cache optimisation. On Apple Silicon we can use a direct replacement, cblas_sgemm from the Accelerate framework. It uses AMX (matrix co-processor) and happens to be 50..100x faster for any matrix above 32×32:

cblas_sgemm(CblasRowMajor,
            ta ? CblasTrans : CblasNoTrans,
            tb ? CblasTrans : CblasNoTrans,
            m, n, k, 1.f,
            A, ta ? m : k,
            B, tb ? k : n,
            0.f, C, n);

But we can go even further and use Metal with GPU matmul for large matrices. Metal comes with MPSMatrix class and a few methods to do the multiplication. But it’s only a 2D matrices, so we’d have to implement our own kernel to handle more than 2D tensors.

kernel void bmatmul(device
  const float * a, device
  const float * b_, device float * c,
    constant int * p, uint idx[[thread_position_in_grid]]) {
  int M = p[1], N = p[2], K = p[3];
  int tot = p[0] * M * N;
  if ((int) idx >= tot) return;
  int n = (int) idx % N, t = (int) idx / N, m = t % M, bat = t / M;
  float s = 0;
  int ao = bat * M * K + m * K, bo = bat * K * N + n;
  for (int k = 0; k < K; k++) s += a[ao + k] * b_[bo + k * N];
  c[idx] = s;
}

The overall ut_matmul function would now look like this:

ut_tensor *ut_matmul(ut_tensor *a, ut_tensor *b) {
    // special case: 2D x 2D
    if (a->shape.ndim == 2 && b->shape.ndim == 2) {
        int n = ..., m = ..., k = ...; // dimensions: a: [m, k], b: [k, n]
        if (a->dev == UT_METAL || b->dev == UT_METAL) {
            // GPU matmul
            ut_tensor *out = ut_alloc(2, (int[]){m, n}, UT_METAL);
            ut_mtl_matmul(a, b, out);
            return out;
        } else {
            // CPU matmul
            ut_tensor *out = ut_alloc(2, (int[]){m, n}, UT_CPU);
            cblas_sgemm(CblasRowMajor,
                        CblasNoTrans, CblasNoTrans,
                        m, n, k, 1.f,
                        a->data, k,
                        b->data, n,
                        0.f, out->data, n);
            return out;
        }
    }
    // special case: 3D x 3D (batches)
    if (a->shape.ndim == 3 && b->shape.ndim == 3) {
        int B, n, m, k; // dimensions: a: [B, m, k], b: [B, k, n]
        // similarly, for Metal: call a custom kernel, for CPU - sgemm in a loop for each batch
    }
}

One thing missing is matmul for transposed matrices. Earlier we found a way to “transpose” a matrix without copying the data (only updating the strides), but for multiplication we should handle it in a special way, too. Fortunately, all implementations seem to support transposition with little code changes (ta and tb flags in cblas_sgemm, and a few extra lines in the Metal kernel).

Autograd by hand

At this point we can perform arithmetics on tensors and adding new operations becomes mostly straightforward. It’s a good foundation for building neural networks, but we are still missing one crucial feature – proper differentiation.

When a network “learns” it adjusts its parameters to minimize a loss function. To do this, we need to compute the gradients of the loss function with respect to the parameters. This is where automatic differentiation usually comes in.

In “big” frameworks like PyTorch or Tensorflow a differentiation engine records a computation graph during the forward pass, and then walks it in reverse to compute the gradients. I’m going to do a simpler thing, manual backpropagation. Every layer of a network becomes a struct with “forward” and “backward” functions and gradients are pre-allocated tensors of the same shape as each trainable parameter.

The forward pass takes an input tensor and produces an output tensor. The backward pass takes the gradient of the loss with respect to the output tensor, and computes the gradient of the loss with respect to the input tensor and the parameters. It may not be as user-friendly as autograd, but for most well-known network architectures it’s easier to implement in C.

Let’s start with a Linear layer. It consists of a weights matrix NxM (where N is number of input features and M is number of output features), and optionally a bias vector:

typedef struct ut_linear {
  ut_tensor* weight;  // [in, out]
  ut_tensor* bias;    // [out]
  int nin, nout;
  bool has_bias;
} ut_linear;

Forward pass is quite simple: a matrix multiplication followed by adding a bias vector (later we might consider fusing these two operations into one to speed things up a little):

ut_tensor *ut_linear_forward(ut_linear *l, ut_tensor *x) {
    ut_tensor *out = ut_matmul(x, l->weight);   // [B, in] @ [in, out]
    if (l->bias)
        for (int b = 0; b < out->shape.shape[0]; b++)
            for (int j = 0; j < l->nout; j++)
                out->data[b * l->nout + j] += l->bias->data[j];
    return out;
}

In practice the forward pass would be a more verbose due to metal/non-metal branches, but the logic remains simple. The backward pass is a bit harder, as we need to compute gradients with respect to the input and the parameters.

A great explanation of matmul backpropagation can be found in this article. For my mathematically deprived brain it was easier to read how pytorch implemented it:

# forward(x: Tensor, w: Tensor): Tensor
y = x@w

# backward(dldy: Tensor): Tensor
dldx = dldy @ w.T
dldw = x.T @ dldy

If we include bias, then y=x@w+b and dldb is just dldy. Bias is a plus operation, so gradient is kept as-is. In case of mini-batch training we should sum output gradients across the batch dimension and that would be our bias gradient.

ut_tensor *ut_linear_backward(ut_linear *l,
                                ut_tensor *x,
                                ut_tensor *grad_out,
                                ut_tensor *dW, ut_tensor *db) {
    // dW += x.T @ grad_out
    ut_tensor *xT    = ut_transpose(x, 0, 1);
    ut_tensor *dW_b  = ut_matmul(xT, grad_out);
    for (int i = 0; i < dW->shape.nelem; i++) dW->data[i] += dW_b->data[i];

    // db += colsum(grad_out)
    for (int b = 0; b < grad_out->shape.shape[0]; b++)
        for (int j = 0; j < l->nout; j++)
            db->data[j] += grad_out->data[b * l->nout + j];

    // dx = grad_out @ W.T 
    ut_tensor *WT = ut_transpose(l->weight, 0, 1);
    ut_tensor *dx = ut_matmul(grad_out, WT);

    ut_free(xT); ut_free(dW_b); ut_free(WT);
    return dx;
}

And this is where things can be optimised further. In backward pass we do a materialising transposition of the weight matrix, essentially copying the whole tensor. Even with GPU acceleration it’s still not optimal. We can implement a special case of matmul that takes a transposed matrix as input, and avoid the extra allocation and copy, ut_matmul_t.

Finally, we are at the point where we can implement and train a simple fully-connected neural network.

Loss

To learn something you should be able to tell right from wrong. This is where loss functions come into play. They measure how far off the predictions are from the actual labels. For classification tasks a common loss function is cross-entropy, so that’s where we start from.

A network tries to predict the probability distribution of classes for each input sample. The cross-entropy loss measures the difference between the predicted distribution and the true distribution (usually a one-hot encoded vector).

A loss can be calculated like this:

float cross_entropy_loss(ut_tensor *logits, const int *labels, int B) {
    float total = 0.f;
    int C = logits->shape.shape[1];
    for (int b = 0; b < B; b++) {
        float *row = logits->data + b * C;
        float mx = row[0];
        for (int j = 1; j < C; j++) if (row[j] > mx) mx = row[j];
        float sum = 0.f;
        for (int j = 0; j < C; j++) sum += expf(row[j] - mx);
        total += logf(sum) + mx - row[labels[b]];
    }
    return total / (float)B;
}

We subtract max(row) before exp() to avoid overflow, keeping every exponential below 1, preventing cases like exp(89), which easily overflows float32 range.

So, the result is a number, the lower it is – the better the network is trained. But we also should use this value to compute the gradient of the loss with respect to the logits, which will be then used in backpropagation:

// Backward: d_logits[b][j] = softmax[b][j] - (j == labels[b])
ut_tensor *cross_entropy_backward(ut_tensor *logits,
                                  const int *labels, int B) {
    int C = logits->shape.shape[1];
    ut_tensor *d = ut_like(logits);
    for (int b = 0; b < B; b++) {
        float *row  = logits->data + b * C;
        float *drow = d->data      + b * C;
        float mx = row[0];
        for (int j = 1; j < C; j++) if (row[j] > mx) mx = row[j];
        float sum = 0.f;
        for (int j = 0; j < C; j++) sum += expf(row[j] - mx);
        float ls = logf(sum) + mx;
        for (int j = 0; j < C; j++) drow[j] = expf(row[j] - ls) / (float)B;
        drow[labels[b]] -= 1.f / (float)B;
    }
    return d;
}

While this is working, we can go further and implement a “fused” version of cross-entropy loss – the one that calculates the loss and the gradient in a single pass, which is more efficient. Fused operations is a common optimisation technique in deep learning frameworks, and a proper tensor library uses them a lot.

Optimiser

Speaking of optimisations, we need to implement an optimiser to update the model parameters based on the gradients computed during backpropagation.

The gradients computed during backpropagation tell us which direction makes the loss worse. To improve the model we move parameters into the opposite direction, scaled by some “learning rate” factor. So, the gradients describe which how impactful the parameters are on the less and which way to move them to make loss smaller, while the optimiser decides how to actually update each parameter: adjusting them too little will take longer to train, adjusting too much might miss the “sweet spot” and make the network learning unstable.

The simplest optimiser rule is stochastic gradient descent (SGD): adjust each parameter by a gradient multiplied by a small number (learning rate).

We can slightly improve it by adding a “momentum”. Instead of following the gradient directly, we accumulate a velocity vector that takes into account the previous updates. This smoothes out oscillations and accelerates through flat regions.

Additionally, we can use gradient clipping: if any gradient component exceeds a certain threshold – we clamp it. This prevents a few bad input samples from destroying the rest of our training.

void ut_sgd_step(ut_sgd* o, float clip) {
 for each parameter p, gradient g:
   g = clamp(g, -clip, clip)
   if momentum > 0: v = mom*v - lr*g;  p += v
   else:            p -= lr*g
   zero gradient for next iteration
}

There are better optimisers, such as Adam, but SGD is probably the simplest to implement first.

MNIST

MNIST dataset is the “hello world” of deep learning. It consists of 60,000 training images and 10,000 test images of handwritten digits (0-9), each image is 28×28 pixels in grayscale. The task is to classify each image into one of the 10 digit classes.

Time to test our tensor library on a practical task!

Data format of the dataset is trivial – raw bytes for images and separate single-byte labels. For training we’ll need an array of shuffled indices to use random batches. Otherwise, it’s pretty straightforward:

// Model: 784 -> 128 -> ReLU -> 10 (known to have good results)
ut_linear fc1 = ut_linear_alloc(784, 128, true, UT_CPU);
ut_linear fc2 = ut_linear_alloc(128, 10, true, UT_CPU);

// Optimiser: SGD + momentum
ut_tensor* params[] = {fc1.weight, fc1.bias, fc2.weight, fc2.bias};
ut_sgd opt = ut_sgd_alloc(params, 4, 0.01f, 0.9f);

// Batch of 64, 15 epochs, 60k / 64 = 937 batches
int B = 64, epochs = 15, batches = train.n / B;
ut_tensor* grad_logits = ut_alloc(2, (int[]){B, 10}, UT_CPU);
int* idx = malloc((size_t)train.n * sizeof(int));
for (int i = 0; i < train.n; i++) idx[i] = i;

for (int ep = 0; ep < epochs; ep++) {
  shuffle(idx, train.n);
  for (int bi = 0; bi < batches; bi++) {
    // pack batch from shuffled indices
    float bx[B * 784]; int bl[B];
    for (int i = 0; i < B; i++) {
      int ii = idx[bi * B + i];
      memcpy(bx + i * 784, train.imgs + ii * 784, 784 * sizeof(float));
      bl[i] = train.labels[ii];
    }
    ut_tensor* x = ut_from_data(2, (int[]){B, 784}, bx, UT_CPU);

    // forward
    ut_tensor *h1 = ut_linear_forward(&fc1, x);
    ut_tensor *h1r = ut_relu(h1);
    ut_tensor *logits = ut_linear_forward(&fc2, h1r);
    // loss + grad
    float loss = ut_cross_entropy(logits, bl, grad_logits);
    // backward
    ut_tensor* dh1r = ut_linear_backward(&fc2, h1r, grad_logits, opt.grads[2], opt.grads[3]);
    ut_tensor* dh1 = ut_relu_backward(dh1r, h1);
    ut_linear_backward(&fc1, x, dh1, opt.grads[0], opt.grads[1]);
    ut_sgd_step(&opt, 5.0f);
    ut_free(x); ut_free(h1); ut_free(h1r); ut_free(logits); ut_free(dh1r); ut_free(dh1);
  }
}

Running on a rusty MacBook M1, I get 0.8 seconds per epoch, and within 5 epochs it reaches 99% accuracy on train set and ~97% on test set. For a small dataset like this there is no difference between CPU and GPU in terms of performance, but on larger models GPU would pay off.

For comparison: PyTorch achieves same accuracy, but takes almost 2 seconds per epoch. On other, larger models, I noticed that C is usually 2-4x faster than PyTorch, especially if most of the operations get fused properly.

What’s next?

This is just the bare minimum. For a proper tensor library we’ll need many more layers and operations, but adding them is a trivial (although, boring) process: add CPU implementation, add GPU kernel, make a dispatch wrapper (either call unary/binary dispatch, or a custom one), add backward pass.

For example, adding a gelu (or any other activation layer) this is roughly 20 lines of code.

The primary goal for me was to build some ground layer to train tinyML/edgeAI models on MacBook, and so far I like it. The full library has convolutional layers (for image processing), pooling, batch normalization, dropout, and a few more optimisers and loss functions. It can load/store tensors into raw binary files and performs well on toy datasets, like CIFAR, HAR, speech commands, etc.

I’m looking into adding an OpenCL backend to support more devices with a few more optimisations.

I even managed to train a tiny GPT2-like transformer, just like minGPT.

What’s missing? I definitely miss autograd, manual backpropagation can be tiresome. Maybe I’ll add it one day, since we have all backward functions implemented – I could add a “tape” to store all operations and some topology sorting to compute gradients in the right order. But for now, I can live without it.

Also a deliberate limitation of being a single-header, header-only tensor library made me sacrifice support of CUDA, which requires kernels to be stored in .cu files and compiled with nvcc (unless there is some other way, I’m not aware of?). Maybe WebGPU acceleration would be possible, I haven’t looked into it yet, but it’s asynchronous nature and rather immature state makes me doubt.

I like the choice of C, and that the user has full control over all memory operations and optimisations. I’ve also experimented a bit with bindings to other languages. C++ was easy – simply wrap all C APIs into classes and the code suddenly becomes much shorter and more readable due to operator overloading. I tried Go bindings, but CGo overhead makes it a bit harder to justify.

This isn’t a competitor to PyTorch. It’s a half-serious anti-PyTorch, small enough to skim through and understand, yet practical enough to build and train various models.

The full library is available at https://github.com/zserge/utensil and if this long post made you curious – clone the repo, add a layer, train a model, break and debug it, any contributions and feedback are welcome!

I hope you’ve enjoyed this article. You can follow – and contribute to – on Github, Mastodon, Twitter or subscribe via rss.

Jul 14, 2026

See also:
AI or ain’t: Neural Networks and more.