DNGPU binary multiplication checkpoints (tinygrad)
These are 15 trained diagonal Neural GPUs (DNGPU) from a tinygrad reproduction of Figure 3 of Freivalds and Liepins, "Improving the Neural GPU Architecture for Algorithm Learning" (arXiv 1702.08727).
Each model multiplies binary numbers: it is trained on inputs of length at most 41 (operands up to 20 bits) and tested on inputs of length 401 (200-bit operands), so the test measures generalization to inputs ten times longer than any seen in training.
The port follows the official TF 1.0 code, LUMII-Syslab/DNGPU@ecc94f7, line by line with task = "bmul" and m = 96 maps.
The model
The input is a sequence of symbols: the bits of a, a separator, and the bits of b, little-endian, with bits stored as 1 and 2 and the separator as 12. An embedding maps each position to m = 96 channels, and one DCGRU cell with a single weight set is then applied L times for an input of length L, so the number of steps grows with the input. The DCGRU is a convolutional GRU whose memory path is a fixed depthwise shift (a third of the channels stay in place, a third shift one position one way, a third the other way), with hard tanh and hard sigmoid nonlinearities. A final 1×1 convolution reads out one of three classes per position (pad, bit 0, bit 1), giving the product's bits. The model has 9m² + 20m + 3 = 84,867 parameters.
Files
glorot/seed{0..9}.npz are ten seeds of the default port, which uses TF 1.0's glorot_uniform initializer as the final upstream code does.
tf012/seed{0..4}.npz are five seeds with TF 0.12's uniform_unit_scaling initializer, which upstream used when the figure was made; it changes only the input and output kernel limits.
Each .jsonl next to a checkpoint is that run's log, one line per evaluation every 200 steps.
Every checkpoint is the state after the step-5000 evaluation, not the best checkpoint of its run, so the runs that drifted late (glorot seeds 2 and 5, tf012 seeds 0 and 3) are stored in their drifted state.
extras/glorot/seed{0..9}.log and extras/tf012/seed{0..4}.log are the runs' stdout: the same evaluation lines as the .jsonl, preceded by the data-generation line whose digest identifies the run's generated training and test data (runs with the same seed share it).
extras/time0.log is a 40-step timing run of the default config, and extras/prof0.log is a 6-step run with tinygrad's per-kernel debug output (ANSI colors kept), used to profile the training step.
An .npz holds the parameters (param/<name>), AdaMax's m/<name> and v/<name>, the data-generator and RNG state, and a JSON blob (host_json, stored as uint8) with the run config and schedule state.
Nothing is pickled, so it loads with allow_pickle=False.
Results
Bit accuracy is the share of correct output bits over all 401 positions. "upstream-32" is upstream's own metric, 32 length-401 test examples drawn by its counter at each evaluation; "1024" is all 1024 length-401 test examples of the run's test set.
| checkpoint | upstream-32 at step 800 (%) | first exact upstream-32 | upstream-32 at 5000 (%) | 1024 at 5000 (%) | wrong of 1024 at 5000 | fewest wrong of 1024 |
|---|---|---|---|---|---|---|
| glorot/seed0.npz | 62.5 | 3200 | 98.19 | 99.155 | 40 | 20 at 3800 |
| glorot/seed1.npz | 82.3 | 2200 | 100.00 | 99.998 | 7 | 7 at 5000 |
| glorot/seed2.npz | 54.1 | 3600 | 88.44 | 87.703 | 432 | 20 at 3800 |
| glorot/seed3.npz | 56.2 | 3400 | 99.99 | 99.425 | 36 | 17 at 4400 |
| glorot/seed4.npz | 57.2 | 4000 | 99.98 | 99.994 | 10 | 9 at 4800 |
| glorot/seed5.npz | 61.4 | never | 93.59 | 93.798 | 327 | 26 at 3600 |
| glorot/seed6.npz | 56.4 | 4000 | 98.05 | 97.197 | 113 | 34 at 4600 |
| glorot/seed7.npz | 54.8 | 3800 | 99.88 | 99.765 | 451 | 22 at 3800 |
| glorot/seed8.npz | 58.8 | 2600 | 99.65 | 99.382 | 425 | 22 at 4800 |
| glorot/seed9.npz | 57.3 | 3000 | 99.47 | 99.768 | 43 | 22 at 4400 |
| tf012/seed0.npz | 72.6 | 3400 | 95.72 | 96.531 | 261 | 28 at 3400 |
| tf012/seed1.npz | 67.3 | 3400 | 99.96 | 99.721 | 27 | 15 at 4200 |
| tf012/seed2.npz | 54.1 | 2800 | 99.98 | 99.983 | 38 | 13 at 4200 |
| tf012/seed3.npz | 60.8 | 2800 | 95.04 | 96.414 | 130 | 21 at 3000 |
| tf012/seed4.npz | 57.4 | 3200 | 100.00 | 99.975 | 29 | 13 at 3600 |
glorot/seed1.npz is the best single model: 7 of the 1024 length-401 test products have at least one wrong bit.
No checkpoint is exact on all 1024, consistent with the paper's Figure 6, which also shows a few wrong outputs at length 401.
Against Figure 3, digitized from its gridlines, the paper's 5-run mean lies inside our 10-seed min..max range at every evaluation. Our mean rises about 50 to 100 steps later (60.1% against the paper's 72.6% at step 800) and is about 0.7 points lower after step 2600, because some seeds drift after reaching 100%. The paper's text claims 99% at about 800 steps, but its own Figure 3 reaches 99% only near step 2000, and none of our runs reach 99% by step 800.
Using a checkpoint
evaluate.py reads only the parameters and run config, regenerates the run's test set from its seed, and runs upstream's test graph.
Run it from the repo root on CPU; the first call spends about a minute compiling the 401-step unrolled graph.
DEV=CPU uv run --no-project --python 3.14 --with tinygrad==0.14.0 --with numpy python evaluate.py glorot/seed1.npz
DEV=CPU uv run --no-project --python 3.14 --with tinygrad==0.14.0 --with numpy python evaluate.py glorot/seed1.npz --multiply 12345 6789
the first prints bit accuracy and wrong examples over the 1024 length-401 test examples, and the second prints the model's product of two numbers.
tinygrad's CPU backend runs this on one core, so the full 1024-example eval takes about 25 minutes per checkpoint (32 batches of about 1.7e12 FLOPs each); --examples 32 takes about a minute.
For glorot/seed0.npz and tf012/seed0.npz it reproduces the logged step-5000 values exactly (99.155% with 40 wrong, 96.531% with 261 wrong).
approaches/dngpu/dngpu.py and data.py are byte-identical to the training code at the revision that wrote the tf012/ checkpoints, so those resume, for example toward the paper's 20000 steps, with python -m approaches.dngpu.dngpu --resume --init tf012 --seed 0 --ckpt tf012/seed0.npz --out tf012/seed0.jsonl --steps 20000 (training needs a CUDA GPU in practice).
The glorot/ checkpoints were written by an earlier revision that differs only in the --init option and the config keys it stores, so their forward pass is identical but --resume rejects their config.
Training
Each run is 5000 steps of upstream's schedule on one RTX 6000 Ada, about 23 minutes: one batch of 32 per length bin [8, 12, 16, 20, 25, 28, 31, 36, 41] per step, dropout keep 0.9, saturation cost with an adaptive weight, AdaMax with gradient clipping, gradient noise, lr 0.005 with decay on plateau.
The paper trains for 20000 steps; Figure 3's sub-5000-step range is the part reproduced here.
The port was checked against upstream run under TF 2.15 tf.compat.v1: data and batches are bit-identical, and forward pass, gradients and 19 AdaMax steps agree to float32 rounding.
TF's random streams (init, dropout masks, gradient noise) cannot be matched, and upstream's fused train step is racy (some gradients read weights AdaMax has already updated); the port uses pre-update weights throughout.
DEVIATIONS and PAPER_VS_CODE in dngpu.py list every difference from upstream and from the paper, the largest paper-versus-code one being the learning rate, which the paper prints as 0.05 while the code uses 0.01 · 48 / m = 0.005.
License
MIT, inheriting the license of LUMII-Syslab/DNGPU; LICENSE keeps Syslab's notice.