/// script
requires-python = ">=3.12"
dependencies = [
"xarray_sql",
"xarray",
"numpy",
"s3fs",
"zarr<3",
]
[tool.uv.sources]
xarray_sql = { path = "..", editable = true }
///
from future import annotations
from typing import Callable
import numpy as np
import xarray as xr
import datetime
import xarray_sql as xql
SIDE=28# images are 28x28; flatten index is height * SIDE + width
WIDTHS= (
SIDE*SIDE,
) # 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.
SKIP_ZERO_PIXELS=True
def fashion_mnist():
"""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.
try:
ds=xr.open_dataset(
"s3://carbonplan-share/xbatcher/fashion-mnist-train.zarr",
engine="zarr",
chunks=None,
backend_kwargs={"storage_options": {"anon": True}},
if"channel"in ds.dims:
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):
images=images/255.0
ds=ds.assign(images=images, labels=ds["labels"].astype("int64"))
except Exception:
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)
n=3*N_SAMPLES
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))
ds=xr.Dataset(
"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],
widths=WIDTHS,
) ->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,).
weights= {
f"layer_{i}": ((f"inp_{i}", f"out_{i}"), init_weight(inp, out))
for i, (inp, out) in enumerate(zip(widths[:-1], widths[1:]))
biases= {
f"bias_{i}": ((f"out_{i}",), init_bias(out))
for i, out in enumerate(widths[1:])
coords= {}
coords.update(
{f"inp_{i}": np.arange(inp) for i, inp in enumerate(widths[:-1])}
coords.update(
{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}"
return ds, names
def main():
rng=np.random.default_rng(1)
mnist=fashion_mnist()
ctx=xql.XarrayContext()
One Dataset splits into two tables: pixels (sample, height, width) and
labels (sample). The dim names are the join keys.
ctx.from_dataset(
"mnist",
mnist,
chunks=dict(sample=CHUNK),
table_names={
("sample", "height", "width"): "pixels",
("sample",): "labels",
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.
data=ctx.sql(f"""
SELECT sample, labels,
CASE WHEN random() <{TRAIN_FRAC} THEN 'train' ELSE 'test' END AS split
FROM mnist.labels
ORDER BY random()
LIMIT {N_SAMPLES}
""").cache()
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.
pixels=ctx.sql("""
SELECT p.sample, p.height, p.width, p.images
FROM mnist.pixels p JOIN data d ON p.sample = d.sample
""").cache()
ctx.register_table("pixels", pixels)
FROM bias b LEFT JOIN gb g
ON b.layer = g.layer AND b.out = g.out
""").cache()
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).
loss=ctx.sql(f"""
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
WHERE e.out = y.labels
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.
acc= (
ctx.sql(f"""
WITH pred AS (
SELECT sample, out,
ROW_NUMBER() OVER (PARTITION BY sample ORDER BY z DESC) AS rk
FROM logits)
SELECT d.split,
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
WHERE p.rk = 1
GROUP BY d.split
.to_pandas()
.set_index("split")["acc"]
print(
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.
trained=xr.Dataset(
f"layer_{i}": ctx.sql(
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"bias_{i}": ctx.sql(
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)}.")
print(trained)
trained.to_zarr(
f"fashion_mnist_mlp_"
f"{datetime.datetime.now().isoformat(timespec='seconds')}.zarr"
if name =="main":
main()
You can’t perform that action at this time.