/** # ============================================================================= # COPYRIGHT © 2025-2026 Konstantin Vladimirovich Grabko. ALL RIGHTS RESERVED. # CMS Manhattan JiRack Technology — PATENT PENDING # # This code is proprietary. # Personal and non-commercial research use is allowed. # Any commercial use, derivative works for profit, or distribution # requires a paid license and 5% royalty. # # It is Java Developers who likes java # # Unauthorized commercial use is strictly prohibited. # Contact: grabko@cmsmanhattan.com # ============================================================================= */ package com.cbsinc.cms.llm.ml; import ai.djl.Device; import ai.djl.Model; import ai.djl.ndarray.*; import ai.djl.ndarray.index.NDIndex; import ai.djl.ndarray.types.DataType; import ai.djl.ndarray.types.Shape; import ai.djl.nn.*; import ai.djl.nn.core.Linear; import ai.djl.training.GradientCollector; import ai.djl.training.ParameterStore; import ai.djl.training.initializer.ConstantInitializer; import ai.djl.training.initializer.NormalInitializer; import ai.djl.util.Pair; import ai.djl.util.PairList; import java.nio.file.Path; import java.nio.file.Paths; /** * JiRack Ternary Transformer — DJL 0.36 port of JiRackTernaryPyTorch_10b_fixed.py. * Same logic, same fixes: * * [FIX-1] Both weights and activations are DEQUANTIZED before the matmul, * so there is no post-matmul rescale anywhere. * [FIX-2] Activation quantization is blended through lambda too — the whole * forward pass is continuous in lambda (smooth QAT warmup; * lambda = 0 is EXACTLY a plain linear). * [FIX-3] BitLinear does NO normalization of its own. The block's RMSNorm * owns normalization; BitLinear does canonical BitNet b1.58 * absmax activation quant + absmean ternary weight quant. * [FIX-4] lambda is stored as a frozen (1,)-shaped Parameter, so it is * serialized in the .params file and survives checkpoint resume. * [FIX-5] RMSNorm computes statistics in FP32 and casts back (BF16-safe). * [FIX-6] GQA: K/V heads are broadcast to query heads (DJL has no SDPA * with enable_gqa, so the materializing fallback path is used). * [FIX-7] RoPE cos/sin tables are precomputed ONCE (interleaved layout, * already repeat_interleave(2)-expanded), shared by all layers. * [FIX-8] GPT-2-style init: residual output projections (out_proj, ffn_w2) * scaled by 1/sqrt(2 * n_layers). Applied post-initialize. * [FIX-9] Ternary export helper: codes in {-1, 0, +1} as INT8 plus a * per-tensor FP32 gamma — the artifact a real ternary inference * kernel would consume. * * NOTE on RoPE layout: this port uses INTERLEAVED rotation (GPT-NeoX / * llama.cpp "type 0"), matching the PyTorch file. HuggingFace Llama * checkpoints use HALF-SPLIT rotation — if you ever init from HF weights, * permute q/k projections first or you will get garbage. * * Engine note: requires the PyTorch engine (stopGradient, stepped slicing, * stack). Run main() for the smoke tests before any real training. */ public class JiRackJDLTernary_10b { // ===================== CONFIG (mutable for smoke mode) ===================== public static int VOCAB_SIZE = 128256; public static int MODEL_DIM = 4096; // hidden_size public static int FFN_HIDDEN_DIM = 18432; // intermediate_size (JiRack 10B) public static int NUM_LAYERS = 32; public static int NUM_HEADS = 32; public static int NUM_KV_HEADS = 8; public static int HEAD_DIM = 128; public static int MAX_SEQ_LEN = 8192; public static final float RMS_EPS = 1e-5f; public static float ROPE_THETA = 500000.0f; public static final float ROPE_SCALE_FACTOR = 1.0f; public static final float INIT_STD = 0.02f; public static final float QUANT_EPS = 1e-5f; private static final Device DEVICE = Device.cpu(); /** Tiny configuration for graph verification / smoke tests. */ public static void smokeConfig() { VOCAB_SIZE = 256; MODEL_DIM = 64; FFN_HIDDEN_DIM = 128; NUM_LAYERS = 2; NUM_HEADS = 4; NUM_KV_HEADS = 2; HEAD_DIM = 16; MAX_SEQ_LEN = 64; RopeCache.reset(); } /* ============================ Model builder ============================ */ public static Block buildModel() { SequentialBlock model = new SequentialBlock(); model.add(new CustomEmbedding(VOCAB_SIZE, MODEL_DIM)); model.add(new FrozenSignatureLayer()); for (int i = 0; i < NUM_LAYERS; i++) { model.add(new TransformerBlock()); } model.add(new RMSNorm(MODEL_DIM)); model.add(Linear.builder().setUnits(VOCAB_SIZE).optBias(false).build()); // lm_head return model; } /** * [FIX-8] GPT-2-style depth-scaled init. Call AFTER block.initialize(). * All weights are already N(0, INIT_STD) via NormalInitializer; here the * residual-stream projections (out_proj, ffn_w2) are rescaled to * INIT_STD / sqrt(2 * NUM_LAYERS). */ public static void applyDepthScaledInit(Block root) { float factor = (float) (1.0 / Math.sqrt(2.0 * NUM_LAYERS)); for (Pair p : root.getParameters()) { String name = p.getKey(); if (name.contains("out_proj") || name.contains("ffn_w2")) { p.getValue().getArray().muli(factor); } } } /** Quantization warmup hook: 0.0 = full precision, 1.0 = full fake-quant. */ public static void setLambda(Block root, float value) { forEachBitLinear(root, bl -> bl.lambdaArray().set(new NDIndex(":"), value)); } public static float getLambda(Block root) { final float[] out = {0.0f}; final boolean[] found = {false}; forEachBitLinear(root, bl -> { if (!found[0]) { out[0] = bl.lambdaArray().getFloat(0); found[0] = true; } }); return out[0]; } private static void forEachBitLinear(Block block, java.util.function.Consumer fn) { if (block instanceof BitLinear) { fn.accept((BitLinear) block); } for (Pair child : block.getChildren()) { forEachBitLinear(child.getValue(), fn); } } /* ---------------------- RoPE cache (interleaved) ---------------------- */ /** * [FIX-7] cos/sin of shape (MAX_SEQ_LEN, HEAD_DIM), already expanded for * interleaved rotation (each frequency duplicated into positions 2i, 2i+1). * Built exactly once; all layers slice [0..T). */ static final class RopeCache { private static float[] cosTable; private static float[] sinTable; private static int builtSeqLen = -1; private static int builtHeadDim = -1; static synchronized void reset() { cosTable = null; sinTable = null; builtSeqLen = -1; builtHeadDim = -1; } private static synchronized void ensureBuilt() { if (cosTable != null && builtSeqLen == MAX_SEQ_LEN && builtHeadDim == HEAD_DIM) { return; } int d = HEAD_DIM; int half = d / 2; double[] invFreq = new double[half]; for (int i = 0; i < half; i++) { double f = 1.0 / Math.pow(ROPE_THETA, 2.0 * i / d); if (ROPE_SCALE_FACTOR > 1.0f) { f /= ROPE_SCALE_FACTOR; } invFreq[i] = f; } cosTable = new float[MAX_SEQ_LEN * d]; sinTable = new float[MAX_SEQ_LEN * d]; for (int t = 0; t < MAX_SEQ_LEN; t++) { for (int i = 0; i < half; i++) { double angle = t * invFreq[i]; float c = (float) Math.cos(angle); float s = (float) Math.sin(angle); // Interleaved expansion: repeat_interleave(2) cosTable[t * d + 2 * i] = c; cosTable[t * d + 2 * i + 1] = c; sinTable[t * d + 2 * i] = s; sinTable[t * d + 2 * i + 1] = s; } } builtSeqLen = MAX_SEQ_LEN; builtHeadDim = d; } /** Returns {cos, sin}, each shaped (1, 1, T, HEAD_DIM). */ static NDArray[] get(NDManager m, long T) { ensureBuilt(); if (T > MAX_SEQ_LEN) { throw new IllegalArgumentException( "Sequence length " + T + " exceeds MAX_SEQ_LEN " + MAX_SEQ_LEN); } int d = HEAD_DIM; int n = (int) T * d; float[] c = new float[n]; float[] s = new float[n]; System.arraycopy(cosTable, 0, c, 0, n); System.arraycopy(sinTable, 0, s, 0, n); return new NDArray[]{ m.create(c, new Shape(1, 1, T, d)), m.create(s, new Shape(1, 1, T, d)) }; } } /* -------------------- BitLinear (b1.58 + lambda warmup) -------------------- */ /** * BitNet b1.58-style fake-quant linear with lambda warmup. * * Forward at lambda = 1: * gamma = mean(|W|) (per-tensor weight scale) * W_q = clamp(round(W / gamma), -1, 1) * gamma // ternary, dequantized * s = 127 / absmax(x, last dim) (per-token activation scale) * x_q = clamp(round(x * s), -128, 127) / s // int8, dequantized * out = x_q @ W_q^T * * Both operands dequantized -> no output rescale [FIX-1]. * STE blend applied to BOTH weights and activations [FIX-2]: * eff = full + lambda * stopGradient(quant - full) * At lambda = 0 this is exactly a plain linear. * No normalization inside [FIX-3] — the block's RMSNorm owns that. */ static class BitLinear extends AbstractBlock { private static final byte VERSION = 1; private final Parameter weight; // (out, in) private final Parameter lambda; // (1,) frozen buffer [FIX-4] private final int inFeatures; private final int outFeatures; public BitLinear(int inFeatures, int outFeatures) { super(VERSION); this.inFeatures = inFeatures; this.outFeatures = outFeatures; weight = addParameter(Parameter.builder() .setName("weight") .setType(Parameter.Type.WEIGHT) .optShape(new Shape(outFeatures, inFeatures)) .build()); lambda = addParameter(Parameter.builder() .setName("lambda") .setType(Parameter.Type.OTHER) .optShape(new Shape(1)) .optInitializer(new ConstantInitializer(0.0f)) .optRequiresGrad(false) // buffer, not trainable .build()); } NDArray lambdaArray() { return lambda.getArray(); } NDArray weightArray() { return weight.getArray(); } @Override public void initializeChildBlocks(NDManager manager, DataType dataType, Shape... inputShapes) {} @Override protected NDList forwardInternal(ParameterStore ps, NDList inputs, boolean training, PairList params) { NDArray x = inputs.get(0); NDArray w = ps.getValue(weight, x.getDevice(), training); NDArray lam = ps.getValue(lambda, x.getDevice(), false); // Fast path: at lambda = 0 the blend is exactly identity, so // skipping the quant math is a pure compute optimization. if (!training && lam.getFloat(0) < 1e-6f) { return new NDList(x.matMul(w.transpose(1, 0))); } // === Weights: per-tensor absmean ternary (stats in FP32, [FIX-5]) === NDArray gamma = w.toType(DataType.FLOAT32, true).abs().mean() .maximum(QUANT_EPS).toType(w.getDataType(), false); NDArray wQuant = w.div(gamma).round().clip(-1.0, 1.0).mul(gamma); NDArray wEff = w.add(wQuant.sub(w).stopGradient().mul(lam)); // === Activations: per-token absmax int8 [FIX-3: no norm here] === NDArray xScale = x.abs().max(new int[]{-1}, true).maximum(QUANT_EPS); NDArray s = xScale.div(127.0f); // = 1 / (127/absmax) NDArray xQuant = x.div(s).round().clip(-128.0, 127.0).mul(s); NDArray xEff = x.add(xQuant.sub(x).stopGradient().mul(lam)); // [FIX-1] Dequantized operands -> plain matmul, no rescale. return new NDList(xEff.matMul(wEff.transpose(1, 0))); } @Override public Shape[] getOutputShapes(Shape[] in) { Shape s = in[0]; long[] dims = s.getShape().clone(); dims[dims.length - 1] = outFeatures; return new Shape[]{new Shape(dims)}; } } /* -------------------------- RMSNorm (FP32 stats) -------------------------- */ static class RMSNorm extends AbstractBlock { private static final byte VERSION = 2; private final Parameter weight; public RMSNorm(int dim) { super(VERSION); weight = addParameter(Parameter.builder() .setName("weight") .setType(Parameter.Type.GAMMA) .optShape(new Shape(dim)) .optInitializer(new ConstantInitializer(1.0f)) .build()); } @Override protected NDList forwardInternal(ParameterStore ps, NDList inputs, boolean training, PairList params) { NDArray x = inputs.get(0); DataType dtype = x.getDataType(); NDArray gamma = ps.getValue(weight, x.getDevice(), training); // [FIX-5] statistics in FP32, cast back NDArray x32 = x.toType(DataType.FLOAT32, false); NDArray normed = x32.mul( x32.pow(2).mean(new int[]{-1}, true).add(RMS_EPS).rsqrt()); return new NDList(normed.mul(gamma.toType(DataType.FLOAT32, false)).toType(dtype, false)); } @Override public void initializeChildBlocks(NDManager manager, DataType dataType, Shape... inputShapes) {} @Override public Shape[] getOutputShapes(Shape[] in) { return in; } } /* --------------------------- TransformerBlock --------------------------- */ static class TransformerBlock extends AbstractBlock { private static final byte VERSION = 3; private final RMSNorm norm1, norm2; private final BitLinear qProj, kProj, vProj, outProj; private final BitLinear ffnW1, ffnW3, ffnW2; private final int h = NUM_HEADS; private final int kvH = NUM_KV_HEADS; private final int d = HEAD_DIM; public TransformerBlock() { super(VERSION); norm1 = addChildBlock("norm1", new RMSNorm(MODEL_DIM)); norm2 = addChildBlock("norm2", new RMSNorm(MODEL_DIM)); qProj = addChildBlock("q_proj", new BitLinear(MODEL_DIM, h * d)); kProj = addChildBlock("k_proj", new BitLinear(MODEL_DIM, kvH * d)); vProj = addChildBlock("v_proj", new BitLinear(MODEL_DIM, kvH * d)); outProj = addChildBlock("out_proj", new BitLinear(h * d, MODEL_DIM)); ffnW1 = addChildBlock("ffn_w1", new BitLinear(MODEL_DIM, FFN_HIDDEN_DIM)); ffnW3 = addChildBlock("ffn_w3", new BitLinear(MODEL_DIM, FFN_HIDDEN_DIM)); ffnW2 = addChildBlock("ffn_w2", new BitLinear(FFN_HIDDEN_DIM, MODEL_DIM)); } @Override public void initializeChildBlocks(NDManager manager, DataType dataType, Shape... inputShapes) { Shape hidden = inputShapes[0]; norm1.initialize(manager, dataType, hidden); norm2.initialize(manager, dataType, hidden); qProj.initialize(manager, dataType, hidden); kProj.initialize(manager, dataType, hidden); vProj.initialize(manager, dataType, hidden); outProj.initialize(manager, dataType, hidden); ffnW1.initialize(manager, dataType, hidden); ffnW3.initialize(manager, dataType, hidden); ffnW2.initialize(manager, dataType, ffnW1.getOutputShapes(new Shape[]{hidden})); } /** Interleaved (NeoX-style) rotation: pairs (x0,x1),(x2,x3)... */ private NDArray applyRotaryInterleaved(NDArray x, NDArray cos, NDArray sin) { long dHead = x.getShape().get(3); NDArray xEven = x.get(new NDIndex("..., 0:" + dHead + ":2")); NDArray xOdd = x.get(new NDIndex("..., 1:" + dHead + ":2")); // rotate = interleave(-x_odd, x_even) NDArray rotated = NDArrays.stack(new NDList(xOdd.neg(), xEven), -1) .reshape(x.getShape()); return x.mul(cos).add(rotated.mul(sin)); } @Override protected NDList forwardInternal(ParameterStore ps, NDList inputs, boolean training, PairList params) { NDArray x = inputs.get(0); NDManager m = x.getManager(); long B = x.getShape().get(0), T = x.getShape().get(1); // Pre-norm attention NDArray hIn = norm1.forward(ps, new NDList(x), training).singletonOrThrow(); NDArray Q = qProj.forward(ps, new NDList(hIn), training).singletonOrThrow() .reshape(B, T, h, d).transpose(0, 2, 1, 3); NDArray K = kProj.forward(ps, new NDList(hIn), training).singletonOrThrow() .reshape(B, T, kvH, d).transpose(0, 2, 1, 3); NDArray V = vProj.forward(ps, new NDList(hIn), training).singletonOrThrow() .reshape(B, T, kvH, d).transpose(0, 2, 1, 3); NDArray[] cs = RopeCache.get(m, T); Q = applyRotaryInterleaved(Q, cs[0], cs[1]); K = applyRotaryInterleaved(K, cs[0], cs[1]); // [FIX-6] GQA fallback path: materialize K/V to query-head count int groups = h / kvH; if (groups > 1) { K = K.reshape(B, kvH, 1, T, d) .broadcast(new Shape(B, kvH, groups, T, d)).reshape(B, h, T, d); V = V.reshape(B, kvH, 1, T, d) .broadcast(new Shape(B, kvH, groups, T, d)).reshape(B, h, T, d); } float scale = (float) (1.0 / Math.sqrt(d)); NDArray scores = Q.matMul(K.transpose(0, 1, 3, 2)).mul(scale); NDArray mask = m.arange(T).reshape(-1, 1).lt(m.arange(T).reshape(1, -1)) .reshape(1, 1, T, T).broadcast(new Shape(B, h, T, T)); NDArray infMask = m.zeros(scores.getShape()).sub(1e9f); scores = NDArrays.where(mask, infMask, scores); NDArray attnOut = scores.softmax(-1).matMul(V) .transpose(0, 2, 1, 3).reshape(B, T, (long) h * d); x = x.add(outProj.forward(ps, new NDList(attnOut), training).singletonOrThrow()); // Pre-norm SwiGLU FFN NDArray mIn = norm2.forward(ps, new NDList(x), training).singletonOrThrow(); NDArray gate = ffnW1.forward(ps, new NDList(mIn), training).singletonOrThrow(); gate = gate.mul(gate.getNDArrayInternal().sigmoid()); // SiLU NDArray up = ffnW3.forward(ps, new NDList(mIn), training).singletonOrThrow(); NDArray down = ffnW2.forward(ps, new NDList(gate.mul(up)), training).singletonOrThrow(); x = x.add(down); return new NDList(x); } @Override public Shape[] getOutputShapes(Shape[] in) { return new Shape[]{in[0]}; } } /* -------------------------- Embedding & signature -------------------------- */ static class CustomEmbedding extends AbstractBlock { private static final byte VERSION = 1; private final Parameter weight; private final int embedDim; public CustomEmbedding(int vocabSize, int embedDim) { super(VERSION); this.embedDim = embedDim; weight = addParameter(Parameter.builder() .setName("embedding_weight") .setType(Parameter.Type.WEIGHT) .optShape(new Shape(vocabSize, embedDim)) .build()); } @Override protected NDList forwardInternal(ParameterStore ps, NDList inputs, boolean training, PairList params) { NDArray ids = inputs.get(0); // integer token IDs (B, T) return new NDList(ps.getValue(weight, ids.getDevice(), training).get(ids)); } @Override public Shape[] getOutputShapes(Shape[] in) { return new Shape[]{new Shape(in[0].get(0), in[0].get(1), embedDim)}; } } static class FrozenSignatureLayer extends AbstractBlock { private static final byte VERSION = 1; @SuppressWarnings("unused") private final Parameter signatureWeight; public FrozenSignatureLayer() { super(VERSION); signatureWeight = addParameter(Parameter.builder() .setName("H32_L32_D4096_V128256_ML8192_cmsmanhattan_model_signature") .setType(Parameter.Type.OTHER) .optShape(new Shape(1, 1)) .optInitializer(new ConstantInitializer(0.0f)) .optRequiresGrad(false) .build()); } @Override public void initializeChildBlocks(NDManager manager, DataType dataType, Shape... inputShapes) {} @Override protected NDList forwardInternal(ParameterStore ps, NDList inputs, boolean training, PairList params) { return inputs; } @Override public Shape[] getOutputShapes(Shape[] in) { return in; } } /* ------------------------ [FIX-9] Ternary export ------------------------ */ /** * Turns trained fake-quant weights into an actual ternary artifact: * for every BitLinear, ".codes" (INT8 in {-1,0,+1}) and * ".gamma" (scalar FP32). Full-precision parts (embedding, * lm_head, norms) are copied as-is. Returned NDList has named arrays * and can be persisted with NDList.encode(). */ public static NDList exportTernaryStateDict(Block root, NDManager manager) { NDList out = new NDList(); exportWalk(root, "", out, manager); for (Pair p : root.getParameters()) { String name = p.getKey(); // Everything not owned by a BitLinear stays full precision. if (!name.contains("weight") || name.contains("embedding") || name.contains("norm") || name.contains("signature") || isLmHead(name)) { NDArray copy = p.getValue().getArray().duplicate(); copy.setName(name); out.add(copy); } } return out; } private static boolean isLmHead(String paramName) { // The final Linear in the SequentialBlock (lm_head) — its weight // must stay FP32; BitLinear weights are exported as codes instead. return paramName.endsWith("Linear.weight"); } private static void exportWalk(Block block, String prefix, NDList out, NDManager manager) { if (block instanceof BitLinear) { BitLinear bl = (BitLinear) block; NDArray w = bl.weightArray().toType(DataType.FLOAT32, true); NDArray gamma = w.abs().mean().maximum(QUANT_EPS); NDArray codes = w.div(gamma).round().clip(-1.0, 1.0).toType(DataType.INT8, false); codes.setName(prefix + "codes"); gamma.setName(prefix + "gamma"); out.add(codes); out.add(gamma); return; } for (Pair child : block.getChildren()) { exportWalk(child.getValue(), prefix + child.getKey() + ".", out, manager); } } /* ------------------------------- Smoke tests ------------------------------- */ public static void main(String[] args) throws Exception { System.out.println("SMOKE MODE: tiny config, mirrors the PyTorch smoke tests"); smokeConfig(); try (Model model = Model.newInstance("jirack", DEVICE); NDManager manager = NDManager.newBaseManager()) { Block block = buildModel(); model.setBlock(block); block.setInitializer(new NormalInitializer(INIT_STD), Parameter.Type.WEIGHT); block.initialize(manager, DataType.FLOAT32, new Shape(2, 32)); applyDepthScaledInit(block); // [FIX-8] ParameterStore psEval = new ParameterStore(manager, false); NDArray ids = manager.randomInteger(0, VOCAB_SIZE, new Shape(2, 32), DataType.INT64); // 1) Continuity in lambda: tiny lambda must NOT jump. setLambda(block, 0.0f); NDArray y0 = block.forward(psEval, new NDList(ids), false).singletonOrThrow(); setLambda(block, 1e-4f); NDArray yEps = block.forward(psEval, new NDList(ids), false).singletonOrThrow(); setLambda(block, 1.0f); NDArray y1 = block.forward(psEval, new NDList(ids), false).singletonOrThrow(); float relJump = yEps.sub(y0).norm().getFloat() / y0.norm().getFloat(); System.out.printf("relative change at lambda=1e-4: %.2e (must be ~1e-4, not ~1)%n", relJump); if (relJump >= 1e-2f) throw new AssertionError("lambda warmup is not continuous!"); // 2) Output scale sanity at full quantization. float ratio = std(y1) / std(y0); System.out.printf("std ratio lambda=1 vs lambda=0: %.3f (must be O(1), not ~1e-4)%n", ratio); if (!(ratio > 0.1f && ratio < 10.0f)) { throw new AssertionError("output scale collapsed or exploded!"); } // 3) Gradients flow through the STE at lambda = 1. setLambda(block, 1.0f); ParameterStore psTrain = new ParameterStore(manager, false); try (GradientCollector gc = manager.getEngine().newGradientCollector()) { NDArray logits = block.forward(psTrain, new NDList(ids), true).singletonOrThrow(); NDArray loss = logits.pow(2).mean(); gc.backward(loss); } NDArray g = findFirstBitLinear(block).weightArray().getGradient(); if (g == null || !g.isFinite().all().getBoolean() || g.abs().sum().getFloat() == 0f) { throw new AssertionError("no finite nonzero gradient through STE!"); } System.out.printf("grad norm on q_proj at lambda=1: %.4f (nonzero, finite)%n", g.norm().getFloat()); // 4) lambda survives a save/load round-trip [FIX-4]. Path tmp = Paths.get("build/smoke_models"); java.nio.file.Files.createDirectories(tmp); model.save(tmp, "jirack"); try (Model model2 = Model.newInstance("jirack", DEVICE)) { Block block2 = buildModel(); model2.setBlock(block2); model2.load(tmp, "jirack"); float lam2 = getLambda(block2); if (Math.abs(lam2 - 1.0f) > 1e-9) { throw new AssertionError("lambda not serialized! got " + lam2); } System.out.println("lambda serialization: OK"); } // 5) Export produces genuinely ternary codes [FIX-9]. NDList exported = exportTernaryStateDict(block, manager); long ternaryCount = exported.stream() .filter(a -> a.getName() != null && a.getName().endsWith("codes")) .count(); NDArray anyCodes = exported.stream() .filter(a -> a.getName() != null && a.getName().endsWith("codes")) .findFirst().orElseThrow(() -> new AssertionError("no codes exported")); float mn = anyCodes.toType(DataType.FLOAT32, false).min().getFloat(); float mx = anyCodes.toType(DataType.FLOAT32, false).max().getFloat(); if (mn < -1f || mx > 1f) throw new AssertionError("codes out of {-1,0,1}!"); System.out.printf("export: %d ternary tensors, values in [%.0f, %.0f], OK%n", ternaryCount, mn, mx); System.out.println("\nAll smoke tests passed."); } } private static float std(NDArray a) { NDArray f = a.toType(DataType.FLOAT32, false); NDArray mean = f.mean(); return (float) Math.sqrt(f.sub(mean).pow(2).mean().getFloat()); } private static BitLinear findFirstBitLinear(Block root) { final BitLinear[] found = {null}; forEachBitLinear(root, bl -> { if (found[0] == null) found[0] = bl; }); return found[0]; } }