torch-type-nn

A PyTorch port of type-nn: a neural network whose layers are partition functions of an And over Ors, and which grows and prunes its own width, degree and depth while it trains. The design is explained in Train the knowledge.

Status: alpha (0.1.0.dev0). Forward, gradients and the per-Or Adam step agree with the C reference to 1e-12, and both C scaling rules have exact ports on the board (see Validation). Deliberate differences from C are listed in docs/DIFFERENCES.md; the audit that led to the current defaults is in docs/AUDIT.md.

Quick start

import torch
from torch_type_nn import TypeNN, fit

model = TypeNN(in_features=4, out_features=3)     # birth: round(ln(1 + n m)) layers
result = fit(model, X, Y, epochs=250, lr=0.005)   # grows and prunes while it trains
print(model.structure(), model.num_params(), result.counters)

fit reproduces type-nn's per-sample protocol by default (batch_size=1); pass batch_size=64 (or more) for the batched, GPU-friendly path. There is no width, degree or depth to choose.

For your own training loop, drive the structure with a StructureScaler:

from torch_type_nn import TypeNN, TypeAdam, StructureScaler

model = TypeNN(n, m)
opt = TypeAdam(model, lr=lr)
scaler = StructureScaler(model, opt, epochs=E, steps_per_epoch=len(loader))
scaler.begin()
for epoch in range(E):
    for x, t in loader:
        y = model(x)
        loss = 0.5 * (y - t).pow(2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
        scaler.observe(x, y, t)          # residual statistics (and pairs, for rule="bic")
    scaler.epoch_end()                   # grow / prune at the boundary
scaler.end()                             # last prune; no probe survives

Structure learning

One dummy rule on three axes. A probe is an identity that back-prop is free to move; a probe that back-prop moved past theta(T) = lr T^(3/4) (T its age in optimizer steps) is promoted, and a fresh probe takes its place. All scaling up happens early and all scaling down late.

(u = step / total; the band is theta_band = lr S^(3/4), S optimizer steps per epoch.) Growth is gated on an unexplained residual (MSE > Var(t) / N). Depth edits fold the best affine fit between x and F(x) into the next layer, so they are near-exact. A network is born with round(ln(1 + n m)) layers, n inputs and m outputs, each layer's output feeding the next, as dense as an MLP.

Two rules. StructureScaler(..., rule=...):

• "threshold" (default) is the architecture: every decision comes from what back-prop did to the dummy weights; no training pair is stored or re-run. It is the rule of C's type-nn-overfit.
• "bic" adds the Bayesian-information prior of the post: a moved probe is promoted only if it also pays n ln(MSE_without / MSE_with) > k ln n, where MSE_without is measured by resetting the item to its identity and re-running the epoch's training pairs (so the scaler keeps one epoch of (x, t) in memory), and pruning ablates items while the criterion n ln MSE + K ln n allows. It is the rule of C's type-nn.

Width on every junction. By default a width probe also grows on a junction that crosses the depth probe (carried through the probe layer). C blocks width there, which leaves a network born with depth ≤ 2 no width site while a depth probe exists; TypeNNAdapter(model, width_through_depth_probe=False) reproduces C.

Both rules and both width rules are on BOARD.md, including exact ports of the two C models.

Any architecture, any optimizer. The scaler is written against a small protocol (torch_type_nn.protocol.Scalable): the rule lives in StructureScaler, and what an item is (how to insert an identity, measure its distance from it, remove it) comes from an adapter. TypeNN is wrapped in TypeNNAdapter automatically. ScalableMLP (Linear/ReLU stacks) is the second family: MLPAdapter grows and prunes its width and depth. See docs/ADAPTERS.md for writing an adapter. Stock torch.optim optimizers work too: structural edits are reported as Edits with index maps, follow_structure remaps the optimizer state through them, and the scaler attaches keep_invariants so a >= 1 holds after every step.

The layer

import torch
from torch_type_nn import AndOr, TypeAdam

layer = AndOr(in_features=8, out_features=4, ors_per_unit=2)
z = layer(torch.randn(32, 8))          # (32, 4)

Unit k of a layer is one And over its Ors r:

Or_kr = w_kr . x + b_kr                sum type
A_k   = prod_r Or_kr ^ a_kr            product type, a_kr >= 1
z_k   = sign(A_k) ln(1 + |A_k|)        partition function

• Batched and padded. All Ors of a layer live in one (m, R, n) tensor; a unit with fewer than R Ors has identity Ors (w = 0, b = 1, a = 1) in its free slots. The identity is exactly 1, so padding never changes the output, and its gradients are masked so no optimizer moves it. This is what gives type-nn batching and GPU support.
• Exact backward. AndOrFunction computes the gradients in log space, including the one-sided rule at a factor that is exactly 0 (cofactor slope when a = 1, flat when a > 1), where autograd through log|O| would give NaN.
• Per-Or Adam. TypeAdam(model, lr) gives every Or its own step count, so an Or born late gets the same bias correction as one born at the start, and keeps a >= 1. It takes the module, not a parameter list, because structural edits replace the parameter tensors; the Adam moments live in the layer (adam_* buffers) so edits move them with the weights. Non-type-nn parameters in the same model get ordinary Adam. Stock torch.optim optimizers also work on a layer of fixed structure.
• Structural primitives. add_or, drop_or, add_unit, drop_unit, add_input (a zero column: function preserved exactly), drop_input, compact. load_state_dict accepts a checkpoint of a different structure.

Validation

Two independent checks against hadilq/type-nn:

• Number for number. tests/test_crosscheck_c.py compiles the C code (tests/c/crosscheck.c), builds the same ragged network in both, and compares the forward pass, dL/dx, every dL/dw, dL/db and dL/da, and the weights after two per-Or Adam steps. Everything agrees to 1e-12. nix flake check runs it (the flake pins the C source as an input); locally, TYPE_NN_SRC=/path/to/type-nn pytest.
• The board. benchmarks/bench.py ports bench.c exactly: the same files, the same xorshift split and shuffle (checked row for row), per-sample Adam at lr x 0.1, the same epochs and learning rates, and 5 seeds per cell. ref-type-nn and ref-type-nn-overfit are exact configurations of C's two models; benchmarks/board.py compares each with the C board (benchmarks/reference/, deterministic, reproduced digit for digit) and writes every table of BOARD.md from the result files.

nix run .#board-all      # every result on all cores, then BOARD.md's tables

Development

Everything is wired through flake.nix:

nix develop            # python + torch + pytest + build tools (or: direnv allow);
                       # also copies the pinned datasets to benchmarks/data
pytest                 # run the tests from the checkout
nix run .#test         # same, without entering the shell
nix flake check        # build the package, run pytest inside the build (with the
                       # pinned datasets), ruff
nix build              # ./result: the installed package
nix run .#dist         # sdist + wheel in ./dist, checked by twine
nix run .#publish -- --repository testpypi   # upload (TestPyPI first)

The benchmark datasets are pinned (URL + SRI hash) in benchmarks/datasets.json and never committed; see benchmarks/data/README.md.

Without Nix: pip install -e ".[dev]" && python benchmarks/fetch_data.py && pytest.

On a GPU

Every test that takes the device fixture also runs on cuda when a CUDA device is visible; TNN_REQUIRE_CUDA=1 makes a missing GPU an error rather than a skip. One script runs the GPU suite (the whole test-suite with TNN_REQUIRE_CUDA=1, then the board's training loop for every model on the device), in two ways:

nix run .#test-cuda              # on the host, now: no Nix configuration needed
nix build .#cuda-tests -L        # in the build sandbox: needs the `cuda` feature
nix flake check --impure         # adds checks.cuda when /dev/nvidiactl exists
TORCH_TYPE_NN_CUDA=1 nix flake check --impure   # force it (=0: leave it out)
nix develop .#cuda               # torch-bin (CUDA 13.0) shell
nix run .#bench-cuda -- all all --seeds 5 --device cuda

A plain nix flake check is a pure evaluation, which cannot see the host, so it never adds the GPU check. With --impure it adds checks.cuda when a GPU is visible (/dev/nvidiactl) and the Nix daemon enables the cuda system feature (read from /etc/nix/nix.conf, /etc/nix/nix.custom.conf, NIX_CONFIG); a GPU without the feature skips the check with a warning. TORCH_TYPE_NN_CUDA=1 / =0 forces it on / off.

The sandboxed check declares requiredSystemFeatures = [ "cuda" ], so the daemon must advertise the feature and expose the GPU to the build:

NixOS:

programs.nix-required-mounts = {
  enable = true;
  presets.nvidia-gpu.enable = true;   # adds the cuda/gpu/opengl features and
};                                    # mounts the driver + /dev/nvidia* for them
# Setting nix.settings.system-features replaces the default list, so keep the
# features you rely on (nixos-test, benchmark, big-parallel, kvm, ...):
nix.settings.system-features = [ "nixos-test" "benchmark" "big-parallel" "kvm" ];

Other Linux (Nix daemon): find the directory holding the driver's libcuda.so.1 (ldconfig -p | grep libcuda.so.1; e.g. /usr/lib/x86_64-linux-gnu on Debian/Ubuntu, /usr/lib on Arch), then in /etc/nix/nix.conf:

extra-system-features = cuda
extra-sandbox-paths = /dev/nvidia0 /dev/nvidiactl /dev/nvidia-uvm /dev/nvidia-uvm-tools? /run/opengl-driver/lib=/usr/lib/x86_64-linux-gnu

and restart the daemon (sudo systemctl restart nix-daemon). torch-bin looks for the driver in /run/opengl-driver/lib; the target=source form mounts the host's driver directory there inside the sandbox only (a trailing ? marks a path that may be missing). With several GPUs add /dev/nvidia1, ... Check with nix build .#cuda-tests -L.

Publishing

Releases go to PyPI from GitHub Actions with trusted publishing (.github/workflows/publish.yml): no API token is stored anywhere.

One-time setup, on PyPI → Your account → Publishing → Add a new pending publisher → GitHub:

For TestPyPI, the same on test.pypi.org (a separate account) with environment testpypi. GitHub creates both environments on first use; add protection rules (e.g. required reviewers) under Settings → Environments if wanted.

Release:

# 1. set the version in pyproject.toml (e.g. 0.1.0), commit, push
# 2. rehearse: Actions → publish → Run workflow → testpypi
# 3. release: the tag must equal the version
git tag v0.1.0 && git push origin v0.1.0

The workflow builds the sdist and wheel, runs the test-suite against the built wheel (CPU torch), and only then uploads. Locally, nix run .#dist builds and checks the same files; nix run .#publish uploads them with twine and a token.

License

MIT. Derived from type-nn by Hadi Lashkati Ghouchani (MIT).