Chris Collins

Tools · math

Neural network designer

Build a multi-layer neural network from a palette, see every tensor shape and parameter count, draw it, train small ones live in the browser, and generate runnable PyTorch, Keras 3, JAX/Flax or plain NumPy code.

Added · updated 14 October 2026

How to use it

Pick a template (or start from Input), then add layers from the palette: click to insert after the selected layer, or drag onto the list. Drag cards (or use ↑ ↓) to reorder; edit sizes inline.

Design

Every layer shows the shape it produces for one example (no batch axis) and its parameter count. Images are height × width × channels, sequences steps × features. Problems are flagged on the layer that causes them — a Dense layer straight after a convolution, a width that does not split into the attention heads, a residual block whose input and output shapes differ.

  • Count parameters as PyTorch, Keras or Flax. They agree except for recurrent layers: PyTorch’s LSTM and GRU keep two bias vectors per gate, Keras’s LSTM and Flax’s cells one. BatchNorm’s running mean and variance are state, not parameters — Keras’s count_params() includes them.

  • Multiply-adds are per example, forward pass (FLOPs ≈ 2 ×). Dense: inputs × outputs. Conv: output positions × kernel area × in-channels × filters. LSTM/GRU: steps × gates × (inputs
    units) × units. Attention: 4·T·d² for the projections + 2·T²·d for the scores.

  • Residual wraps the next N layers and adds its input back after them (a skip connection).

  • The output head chooses the loss: cross-entropy (multi-class), binary cross-entropy (binary, multi-label) or mean squared error (regression). A regression head whose target is the input makes an autoencoder.

The diagram draws small dense networks as neurons and connections — after training, edges are coloured by weight (blue positive, orange negative) — and anything else as blocks with the tensor shape on each arrow. Download it as SVG or PNG.

Train in the browser

Dense networks with one or two inputs train live on toy datasets: XOR, circles, spirals, three Gaussian blobs, and a noisy sine for regression. The left picture shows the data (filled points train, hollow points test) over the network’s decision regions; the right one plots training and test loss and test accuracy. Try a network that is too small for the spirals, then add a layer; add noise and watch the test loss part from the training loss; raise the learning rate until training diverges. Adapt the network to this dataset fixes the input size and output head.

The trainer is plain JavaScript: mini-batch backpropagation with SGD, momentum or Adam, and L2 regularisation. Convolutional, recurrent and attention networks are designed and exported here, but train them with the generated code.

⛶ Full screen (or F) gives the current tab’s main view the whole screen — the decision boundary and loss plots while training, the network diagram on Design, the code on Code — using the browser’s full screen where it is allowed, otherwise a full-window layer (iPhone); the plots are redrawn sharp at the new size, and ◫ Panels hides the training controls or the layer editor. Escape leaves.

Code

  • PyTorch 2.x — an nn.Module, DataLoader, AdamW and a training loop that picks CUDA, Apple mps or the CPU. Inputs are channels-first (N, C, H, W), as PyTorch expects.

  • Keras 3 — the functional API; runs on TensorFlow, JAX or PyTorch (set KERAS_BACKEND).

  • JAX + Flax NNX + Optax — the current flax.nnx API with a jitted training step.

  • NumPy — dense networks only, with the forward pass, backpropagation and Adam written out by hand; python model_numpy.py --check compares the gradients with finite differences.

Each script runs as is: with the playground’s dataset when the network fits it, otherwise on random placeholder data — replace make_data() with your own. Every template’s code, and a set of harder test networks, was run with PyTorch 2.14, Keras 3.15 and Flax 0.12, and each printed the same parameter count as this page.

Sharing

Share link copies a URL that opens this page with your network (and training settings). Templates have short links: ?template= is one of mlp, xor, spiral, blobs, lenet, resnet, conv1d, lstm, transformer, autoencoder, optionally with &fw=keras (or flax, numpy) and #code to open the Code tab — for example ?template=transformer&fw=flax#code.

Linking to an example

Companion to the neuron playground in chapter 33 of Mathematics, Reconnected. A link can build a small network with two inputs, one or more hidden layers, and an output head that fits the dataset, and open it on the Train tab ready to press Play, for example /tools/neural-net/?data=circles&hidden=4&act=tanh#train.

Parameter Meaning Example

data

xor, circles, spiral, blobs (three classes), or sine (regression, one input)

data=spiral

hidden

neurons in each hidden layer (1–64)

hidden=8

layers

hidden layers (1–4)

layers=2

act

activation of the hidden layers: tanh, relu, sigmoid, gelu, silu, leaky_relu, or linear (none)

act=linear

noise, lr

dataset noise (0–0.5) and learning rate

lr=0.01

template

one of the designer’s templates instead

template=lenet

fw

code framework: torch, keras, flax, numpy

fw=numpy

example

one of the presets below; other parameters override it

example=ch33-spiral&hidden=16

Presets:

  • ch33-circle-in-ring — A circle inside a ring: 2 → 4 tanh → 1 bends the boundary into a loop

  • ch33-no-activation — The same data with no activation: two layers, and the boundary stays a straight line

  • ch33-spiral — Two spirals need more: 2 → 8 → 8 tanh → 1

  • ch33-xor — XOR, the classic a single neuron cannot learn: 2 → 4 tanh → 1

  • ch33-blobs — Three blobs, three classes: 2 → 8 ReLU → 3

Runs entirely in your browser; nothing you type leaves this page.