mixle.engines.lns_nn module¶
Neural-network ops in the logarithmic number system – the log-space parts of NNs as integer math.
Built on mixle.engines.lns.LogNumberSystem. Two families of op benefit from LNS, both because
they fight transcendentals (exp/log), not BLAS:
softmax / cross-entropy / log-softmax – the normalizer is a
logsumexpover logits (the LM head over the vocabulary, a classifier over classes, an attention/MoE-router softmax). The integerlogsumexpreplaces theexp``+``logwith integermax+ a LUT (~2x measured). The model’s logits are quantized by the same log step; only softmax-back-to-linear still needs anexp.sum-product circuits / probabilistic circuits – the whole forward pass is sums and products of probabilities, so in LNS every product node is an integer ADD and every sum node an integer
logadd. The entire network runs in integer log-space, not just the normalizer. This is the class of “neural” model where LNS is a complete fit (and the one mixle’s composable probabilistic models naturally express).
These are inference/scoring ops (the gradient path stays in float autograd); LNS buys speed + a 32x-smaller integer representation where the math is log-space.
- log_softmax(logits, lns, axis=-1)[source]
Log-softmax via the integer log-partition:
(k - logsumexp_int(k)) * step. Stays in log-space.
- softmax(logits, lns, axis=-1)[source]
Softmax with the normalizer computed in LNS; back-to-linear needs one
exp(for attention.V etc.).
- cross_entropy(logits, targets, lns, axis=-1)[source]
Mean negative log-likelihood
mean(logsumexp(logits) - logit[target])via the integer normalizer.The LM / classifier loss: the log-partition over the vocab/classes is an integer
logsumexp; the target logit is gathered from the same quantized logits, so the loss is integer until the final scale.
- class SumProductCircuit(nodes)[source]
Bases:
objectA probabilistic circuit evaluated entirely in integer log-space (product=add, sum=logadd).
nodesis a topologically ordered list (children before parents), each a tuple:("leaf", leaf_id)– an input whose log-value is supplied at evaluation,("product", [child indices])– log-output = sum of children (integer ADD),("sum", [child indices], [log_weights])– log-output = logsumexp of weighted children.
The root is the last node.
evaluate_lnsruns the whole forward pass on integers;evaluate_floatis the float64 reference. Leaf values may be scalars or arrays (the forward broadcasts).- evaluate_lns(lns, leaf_values)[source]