Show HN: I implemented a neural network in SQL

# requires-python = ">=3.12"

# xarray_sql = { path = "..", editable = true }

from __future__ import annotations

from typing import Callable

SIDE = 28 # images are 28x28; flatten index is height * SIDE + width

) # 784 pixels -> 196 -> 32 tanh -> 10 softmax

N_SAMPLES, TRAIN_FRAC = 700, 0.7 # total samples; fraction used for training

LR, STEPS, CHUNK = 0.5, 60, 250

# Drop zero-valued pixels from the (dominant) layer-0 contraction. A background

# pixel contributes 0 * weight = 0, so skipping those rows shrinks the join

# *exactly* — the result is identical, and the speedup scales with the fraction

# of zeros (a dark background). On dense inputs it is a no-op.

# Measured ~1.8x on real Fashion-MNIST (~50% zero pixels): 2.56 -> 1.45 s/step.

"""The whole training set, left lazy so SQL streams and samples it.

The real path returns a dask-backed (chunked) Dataset — nothing is pulled

into memory here; ``from_dataset`` reads it chunk by chunk on demand, and

the random subsample happens later in SQL. The offline fallback is a small

synthetic set built in memory.

"s3://carbonplan-share/xbatcher/fashion-mnist-train.zarr",

backend_kwargs={"storage_options": {"anon": True}},

ds = ds.isel(channel=0, drop=True)

# To float64, lazily (no full read). This zarr already stores images

# as float in [0, 1]; only integer-encoded sources ([0, 255]) rescale.

images = ds["images"].astype("float64")

if not np.issubdtype(ds["images"].dtype, np.floating):

ds = ds.assign(images=images, labels=ds["labels"].astype("int64"))

# Offline fallback: a separable synthetic set (per-class template +

# noise), so the same pipeline still learns without the network. A pool

# larger than N_SAMPLES so the SQL subsample still has something to pick.

rng = np.random.default_rng(0)

templates = rng.standard_normal((10, SIDE, SIDE))

labels = rng.integers(0, 10, n).astype("int64")

images = templates[labels] + 0.6 * rng.standard_normal((n, SIDE, SIDE))

"images": (("sample", "height", "width"), images),

"labels": (("sample",), labels),

# Integer index coords are the SQL join keys (sample, height, width).

return ds[["images", "labels"]].assign_coords(

sample=np.arange(ds.sizes["sample"]),

height=np.arange(ds.sizes["height"]),

width=np.arange(ds.sizes["width"]),

def build_model_with_table_names(

init_weight: Callable[[int, int], np.ndarray],

init_bias: Callable[[int], np.ndarray],

) -> tuple[xr.Dataset, dict[tuple[str, ...], str]]:

"""The network as one Dataset that splits into tables per layer.

Layer ``i`` is a weight matrix ``layer_i (inp_i, out_i)`` and a separate

bias vector ``bias_i (out_i,)``.

f"layer_{i}": ((f"inp_{i}", f"out_{i}"), init_weight(inp, out))

for i, (inp, out) in enumerate(zip(widths[:-1], widths[1:]))

f"bias_{i}": ((f"out_{i}",), init_bias(out))

for i, out in enumerate(widths[1:])

{f"inp_{i}": np.arange(inp) for i, inp in enumerate(widths[:-1])}

{f"out_{i}": np.arange(out) for i, out in enumerate(widths[1:])}

ds = xr.Dataset({**weights, **biases}, coords=coords)

names: dict[tuple[str, ...], str] = {}

for i in range(len(weights)):

names[(f"inp_{i}", f"out_{i}")] = f"layer{i}"

names[(f"out_{i}",)] = f"bias{i}"

rng = np.random.default_rng(1)

ctx = xql.XarrayContext()

# One Dataset splits into two tables: pixels (sample, height, width) and

# labels (sample). The dim names are the join keys.

chunks=dict(sample=CHUNK),

("sample", "height", "width"): "pixels",

# Draw a random N_SAMPLES subset in SQL (ORDER BY random() LIMIT), carrying

# each sample's label and a train/test tag. `data` is the working label

# table: cache() pins the chosen subset so every downstream query sees the

# same split without rescanning the source. `ORDER BY random()` shuffles the

# whole label column, so the subset is order-independent even if the on-disk

# samples are class-sorted.

CASE WHEN random() < {TRAIN_FRAC} THEN 'train' ELSE 'test' END AS split

ctx.register_table("data", data)

# Materialise just the sampled images once: a single lazy scan of the full

# dataset extracts the ~N_SAMPLES subset into `pixels`, which the per-step

# forward joins instead of rescanning the source 60x. Only the subset lives

# in memory; the full set stays lazy.

SELECT p.sample, p.height, p.width, p.images

FROM mnist.pixels p JOIN data d ON p.sample = d.sample

ctx.register_table("pixels", pixels)

# The gradient averages over the actual train count (random, ~frac * N),

# read once from the materialized split.

"SELECT COUNT(*) AS n FROM data WHERE split = 'train'"

def init_weight(inp: int, out: int):

"""Small random weights."""

return rng.standard_normal((inp, out)) * 0.1

"""Biases start at zero."""

model, table_names = build_model_with_table_names(init_weight, init_bias)

# Each layer table is one chunk: weights on (inp_i, out_i) and the bias

# vector on (out_i,), so every dim needs a size here.

f"inp_{i}": model.sizes[f"inp_{i}"]

for i in range(len(WIDTHS) - 1)

f"out_{i}": model.sizes[f"out_{i}"]

for i in range(len(WIDTHS) - 1)

# Unify the per-layer weight tables into one working weight(layer, inp, out,

# val) relation the loop rewrites in place, tagging each layer with its

seed = " UNION ALL ".join(

f"SELECT {i} AS layer, inp_{i} AS inp, out_{i} AS out, layer_{i} AS val "

for i in range(len(WIDTHS) - 1)

ctx.register_table("weight", ctx.sql(seed).cache())

# The biases live in their own bias(layer, out, val) relation, summed into

# each layer's pre-activation as a separate term (z = W @ a + b).

bias_seed = " UNION ALL ".join(

f"SELECT {i} AS layer, out_{i} AS out, bias_{i} AS val FROM model.bias{i}"

for i in range(len(WIDTHS) - 1)

ctx.register_table("bias", ctx.sql(bias_seed).cache())

# The zero-pixel skip. fwd0 has no WHERE (it forwards all samples), so it

# needs a fresh `WHERE`; g0 already filters to the train split, so it

# appends an `AND`. Empty strings when the flag is off.

zero_where = "WHERE images <> 0" if SKIP_ZERO_PIXELS else ""

zero_and = "AND images <> 0" if SKIP_ZERO_PIXELS else ""

for step in range(STEPS):

# --- forward pass -----------------------------------------------------

# Each layer contracts its activation with the weight table (JOIN on the

# shared index + grouped SUM), then adds the layer's bias as a separate

# term (JOIN the bias table on `out`), and keeps the pre-activation z

# (tanh(z) for hidden, linear output). .cache() materialises each stage

# so the per-step plan stays flat.

# The forward runs over ALL samples: train rows drive learning, test

# rows ride along so we can score them from the same logits. Only delta2

# is restricted to train, so the gradients (and the trained weights) are

# identical to a train-only forward — test is never backpropagated.

-- z = x @ W: matmul of the input and first weight matrix

SELECT a.sample, w.out AS out, SUM(a.val * w.val) AS z

SELECT sample, height * {SIDE} + width AS inp, images AS val

JOIN weight w ON a.inp = w.inp AND w.layer = 0

-- activation(z + b): Add in the bias term, then perform activation

SELECT c.sample, c.out AS out, c.z + b.val AS z,

FROM c JOIN bias b ON c.out = b.out AND b.layer = 0

ctx.deregister_table("fwd0")

ctx.register_table("fwd0", fwd0)

SELECT a.sample, w.out AS out, SUM(a.val * w.val) AS z

FROM (SELECT sample, out AS inp, val FROM fwd0) a

JOIN weight w ON a.inp = w.inp AND w.layer = 1

SELECT c.sample, c.out AS out, c.z + b.val AS z,

FROM c JOIN bias b ON c.out = b.out AND b.layer = 1

ctx.deregister_table("fwd1")

ctx.register_table("fwd1", fwd1)

# Output layer is linear (softmax lives in the loss / output error),

# but still gets its bias summed in.

SELECT a.sample, w.out AS out, SUM(a.val * w.val) AS z

FROM (SELECT sample, out AS inp, val FROM fwd1) a

JOIN weight w ON a.inp = w.inp AND w.layer = 2

SELECT c.sample, c.out AS out, c.z + b.val AS z

FROM c JOIN bias b ON c.out = b.out AND b.layer = 2

ctx.deregister_table("logits")

ctx.register_table("logits", logits)

# --- backward pass ----------------------------------------------------

# Output error delta2 = softmax(logits) - onehot(label). The one

# hand-derived rule: softmax couples classes through a per-sample

WITH m AS (SELECT sample, MAX(z) AS m FROM logits GROUP BY sample),

e AS (SELECT logits.sample, logits.out, exp(logits.z - m.m) AS e

FROM logits JOIN m ON logits.sample = m.sample),

s AS (SELECT sample, SUM(e) AS s FROM e GROUP BY sample)

e.e / s.s - CASE WHEN e.out = y.labels THEN 1.0 ELSE 0.0 END AS val

FROM e JOIN s ON e.sample = s.sample

JOIN data y ON y.sample = e.sample

-- restrict the error to train, so every downstream gradient is train-only

WHERE e.sample IN (SELECT sample FROM data WHERE split = 'train')

ctx.deregister_table("delta2")

ctx.register_table("delta2", delta2)

# Weight gradient of layer 2: fwd1.T @ delta2 / N.

SELECT a.inp AS inp, d.out AS out, SUM(a.val * d.val) / {n_train} AS val

FROM (SELECT sample, out AS inp, val FROM fwd1) a

JOIN delta2 d ON a.sample = d.sample

ctx.deregister_table("g2")

ctx.register_table("g2", g2)

# Bias gradient of layer 2: the mean output error per unit.

SELECT out, SUM(val) / {n_train} AS val FROM delta2 GROUP BY out

ctx.deregister_table("gb2")

ctx.register_table("gb2", gb2)

# Propagate to layer 1: delta1 = (delta2 @ W2.T) * tanh'(z1). The local

# derivative is grad(tanh(z), z) at fwd1's pre-activation.

SELECT d.sample, w.inp AS out, SUM(d.val * w.val) AS val

FROM delta2 d JOIN weight w ON d.out = w.out AND w.layer = 2

SELECT dc.sample, dc.out,

dc.val * grad(tanh(fwd1.z), fwd1.z) AS val

FROM dc JOIN fwd1 ON dc.sample = fwd1.sample AND dc.out = fwd1.out

ctx.deregister_table("delta1")

ctx.register_table("delta1", delta1)

SELECT a.inp AS inp, d.out AS out, SUM(a.val * d.val) / {n_train} AS val

FROM (SELECT sample, out AS inp, val FROM fwd0) a

JOIN delta1 d ON a.sample = d.sample

ctx.deregister_table("g1")

ctx.register_table("g1", g1)

SELECT out, SUM(val) / {n_train} AS val FROM delta1 GROUP BY out

ctx.deregister_table("gb1")

ctx.register_table("gb1", gb1)

# Propagate to layer 0: delta0 = (delta1 @ W1.T) * tanh'(z0).

SELECT d.sample, w.inp AS out, SUM(d.val * w.val) AS val

FROM delta1 d JOIN weight w ON d.out = w.out AND w.layer = 1

SELECT dc.sample, dc.out,

dc.val * grad(tanh(fwd0.z), fwd0.z) AS val

FROM dc JOIN fwd0 ON dc.sample = fwd0.sample AND dc.out = fwd0.out

ctx.deregister_table("delta0")

ctx.register_table("delta0", delta0)

SELECT sample, height * {SIDE} + width AS inp, images AS val

WHERE sample IN (SELECT sample FROM data WHERE split = 'train')

SELECT a.inp AS inp, d.out AS out, SUM(a.val * d.val) / {n_train} AS val

FROM a JOIN delta0 d ON a.sample = d.sample

ctx.deregister_table("g0")

ctx.register_table("g0", g0)

SELECT out, SUM(val) / {n_train} AS val FROM delta0 GROUP BY out

ctx.deregister_table("gb0")

ctx.register_table("gb0", gb0)

# --- SGD update: one query per relation -------------------------------

# weight <- weight - lr * gradient and bias <- bias - lr * gradient,

# joining every layer at once against the per-layer gradients tagged

# with their layer index.

SELECT 0 AS layer, inp, out, val FROM g0

UNION ALL SELECT 1 AS layer, inp, out, val FROM g1

UNION ALL SELECT 2 AS layer, inp, out, val FROM g2

SELECT w.layer, w.inp, w.out,

w.val - {LR} * COALESCE(g.val, 0) AS val

FROM weight w LEFT JOIN grad g

ON w.layer = g.layer AND w.inp = g.inp AND w.out = g.out

ctx.deregister_table("weight")

ctx.register_table("weight", w)

SELECT 0 AS layer, out, val FROM gb0

UNION ALL SELECT 1 AS layer, out, val FROM gb1

UNION ALL SELECT 2 AS layer, out, val FROM gb2

b.val - {LR} * COALESCE(g.val, 0) AS val

FROM bias b LEFT JOIN gb g

ON b.layer = g.layer AND b.out = g.out

ctx.deregister_table("bias")

ctx.register_table("bias", b)

if step % 5 == 0 or step == STEPS - 1:

# Train cross-entropy (logits span all samples, so filter to train).

WITH m AS (SELECT sample, MAX(z) AS m FROM logits GROUP BY sample),

e AS (SELECT logits.sample, logits.out, exp(logits.z - m.m) AS e

FROM logits JOIN m ON logits.sample = m.sample),

s AS (SELECT sample, SUM(e) AS s FROM e GROUP BY sample)

SELECT -AVG(ln(e.e / s.s)) AS loss

FROM e JOIN s ON e.sample = s.sample

JOIN data y ON y.sample = e.sample

AND e.sample IN (SELECT sample FROM data WHERE split = 'train')

""").to_pandas()["loss"][0]

# Accuracy per split: argmax the shared logits, join the split label.

# Both come from the one all-samples forward — no second pass.

ROW_NUMBER() OVER (PARTITION BY sample ORDER BY z DESC) AS rk

AVG(CASE WHEN p.out = d.labels THEN 1.0 ELSE 0.0 END) AS acc

FROM pred p JOIN data d ON d.sample = p.sample

.set_index("split")["acc"]

f"step {step:2d}: loss {loss:.3f} "

f"train_acc {acc['train']:.3f} test_acc {acc['test']:.3f}"

# The trained parameters come back out as xarray in the *same shape as the

# input model*: one weight variable per layer with its own (inp_i, out_i)

# dims, plus one bias variable per layer on (out_i,). Each is read from its

# relation by the `layer` column, so the result is a ragged set of per-layer

# matrices and vectors — no dense array padded with NaN.

f"SELECT inp AS inp_{i}, out AS out_{i}, val AS layer_{i} "

f"FROM weight WHERE layer = {i}"

).to_dataset(dims=[f"inp_{i}", f"out_{i}"])[f"layer_{i}"]

for i in range(len(WIDTHS) - 1)

f"SELECT out AS out_{i}, val AS bias_{i} "

f"FROM bias WHERE layer = {i}"

).to_dataset(dims=[f"out_{i}"])[f"bias_{i}"]

for i in range(len(WIDTHS) - 1)

print(f"trained {WIDTHS} MLP; weights -> xarray {dict(trained.sizes)}.")

f"{datetime.datetime.now().isoformat(timespec='seconds')}.zarr"

if __name__ == "__main__":