chore(refactor): small refactoring to names and tensor logic

xEcho1337 · xEcho1337 · commit 4323adb19824 · 2026-01-03T14:59:23.000+01:00
fix(transformer): fixed gate projection throwing NPE when freezing/unfreezing
diff --git a/brain4j-core/src/main/java/org/brain4j/core/Brain4J.java b/brain4j-core/src/main/java/org/brain4j/core/Brain4J.java
@@ -123,7 +123,11 @@ public static Device firstDevice() {
             throw new IllegalStateException("No GPU-acceleration device has been found!");
         }
 
-        return DeviceUtils.findDevice(devices.getFirst());
+        Device device = DeviceUtils.findDevice(devices.getFirst());
+        
+        if (device != null) Brain4J.initKernels(device);
+        
+        return device;
     }
 
     /**
diff --git a/brain4j-core/src/main/java/org/brain4j/core/importing/SafeTensorsConverter.java b/brain4j-core/src/main/java/org/brain4j/core/importing/SafeTensorsConverter.java
@@ -114,6 +114,9 @@ private static Map<String, Tensor> load(ByteBuffer buffer) throws IOException {
         
         for (Map.Entry<String, JsonElement> entry : header.entrySet()) {
             String name = entry.getKey();
+            
+            if (name.equals("__metadata__")) continue;
+            
             JsonObject info = entry.getValue().getAsJsonObject();
             
             JsonArray shapeArray = info.getAsJsonArray("shape");
diff --git a/brain4j-core/src/main/java/org/brain4j/core/layer/impl/RecurrentLayer.java b/brain4j-core/src/main/java/org/brain4j/core/layer/impl/RecurrentLayer.java
@@ -75,7 +75,7 @@ public Tensor[] forward(StatesCache cache, Tensor... inputs) {
         for (int t = 0; t < timesteps; t++) {
             Range[] ranges = new Range[] { Range.all(), Range.point(t), Range.all() };
 
-            Tensor timestepX = projectedInput.sliceGrad(ranges).squeeze(1);
+            Tensor timestepX = projectedInput.sliceGrad(ranges).squeezeGrad(1);
             Tensor timestepH = hiddenState.matmulGrad(hiddenWeights);
 
             hiddenState = timestepX.addGrad(timestepH).addGrad(hiddenBias).activateGrad(activation);
diff --git a/brain4j-core/src/main/java/org/brain4j/core/layer/impl/transformer/MultiHeadAttention.java b/brain4j-core/src/main/java/org/brain4j/core/layer/impl/transformer/MultiHeadAttention.java
@@ -133,9 +133,9 @@ public Tensor[] forward(StatesCache cache, Tensor... inputs) {
             } else {
                 Tensor reshaped = QKV.reshape(batch, seqLength, H, 3, d)
                         .transpose(1, 2); // [B,H,L,3,d]
-                Q = reshaped.slice(all, all, all, Range.point(0), all).squeeze(3);
-                K = reshaped.slice(all, all, all, Range.point(1), all).squeeze(3);
-                V = reshaped.slice(all, all, all, Range.point(2), all).squeeze(3);
+                Q = reshaped.slice(all, all, all, Range.point(0), all).squeezeGrad(3);
+                K = reshaped.slice(all, all, all, Range.point(1), all).squeezeGrad(3);
+                V = reshaped.slice(all, all, all, Range.point(2), all).squeezeGrad(3);
             }
 
             float scale = (float) (1.0 / Math.sqrt(d));
diff --git a/brain4j-core/src/main/java/org/brain4j/core/layer/impl/transformer/TransformerEncoder.java b/brain4j-core/src/main/java/org/brain4j/core/layer/impl/transformer/TransformerEncoder.java
@@ -283,24 +283,26 @@ public void loadWeights(Map<String, Tensor> mappedWeights) {
     @Override
     public Layer freeze() {
         upProjection.freeze();
-        gateProjection.freeze();
-        gateProjection.freeze();
         downProjection.freeze();
         normalizer1.freeze();
         normalizer2.freeze();
         attention.freeze();
+        
+        if (gateProjection != null) gateProjection.freeze();
+        
         return super.freeze();
     }
     
     @Override
     public Layer unfreeze() {
         upProjection.unfreeze();
-        gateProjection.unfreeze();
-        gateProjection.unfreeze();
         downProjection.unfreeze();
         normalizer1.unfreeze();
         normalizer2.unfreeze();
         attention.unfreeze();
+        
+        if (gateProjection != null) gateProjection.unfreeze();
+        
         return super.unfreeze();
     }
     
diff --git a/brain4j-core/src/main/java/org/brain4j/core/model/impl/Sequential.java b/brain4j-core/src/main/java/org/brain4j/core/model/impl/Sequential.java
@@ -59,7 +59,8 @@ public Tensor[] predict(StatesCache cache, Tensor... inputs) {
                 input = input.reshape(1, input.elements()); // reshape to [batch, input_size]
             }
             
-            buffer[i] = cache.isTraining() ? input.withGrad() : input;
+            Tensor chosen = cache.isTraining() ? input.withGrad() : input;
+            buffer[i] = chosen.to(device);
         }
         
         for (Layer layer : layers) {
diff --git a/brain4j-core/src/main/java/org/brain4j/core/transformer/attention/MaskedMultiHeadAttention.java b/brain4j-core/src/main/java/org/brain4j/core/transformer/attention/MaskedMultiHeadAttention.java
@@ -114,7 +114,7 @@ public Tensor[] forward(StatesCache cache, Tensor... inputs) {
         }
 
         Range[] slicingRanges = {
-                Range.all(), Range.point(seqLength - 1), Range.all()
+            Range.all(), Range.point(seqLength - 1), Range.all()
         }; // [batch, 1, dim]
         Tensor cachedOutput = cache.get(outProj);
         Tensor cachedQKV = cache.get(weights);
diff --git a/brain4j-core/src/main/java/org/brain4j/core/transformer/attention/head/FlashAttentionHead.java b/brain4j-core/src/main/java/org/brain4j/core/transformer/attention/head/FlashAttentionHead.java
@@ -66,11 +66,8 @@ public Tensor attend(Tensor input) {
             }
 
             if (context != null) {
-                return training
-                    ? context.squeezeGrad(1) // [B,L,d]
-                    : context.squeeze(1);
+                return context.squeezeGrad(1); // [B,L,d]
             }
-            // fallthrough if context null
         }
 
         // fallback to standard path with autograd support
diff --git a/brain4j-llm/src/main/java/org/brain4j/llm/core/architecture/impl/GPT2Adapter.java b/brain4j-llm/src/main/java/org/brain4j/llm/core/architecture/impl/GPT2Adapter.java
@@ -10,6 +10,7 @@
 import org.brain4j.core.layer.impl.transformer.TransformerDecoder;
 import org.brain4j.core.layer.impl.utility.InputLayer;
 import org.brain4j.core.model.Model;
+import org.brain4j.core.model.ModelSpecs;
 import org.brain4j.llm.core.architecture.ArchitectureAdapter;
 import org.brain4j.math.data.StatesCache;
 import org.brain4j.math.tensor.Tensor;
@@ -31,7 +32,7 @@ public Model buildModel(JsonObject config, Map<String, Tensor> weights) {
         int context = config.get("n_ctx").getAsInt();
         int vocabSize = config.get("vocab_size").getAsInt();
         
-        OldSequential seq = OldSequential.of();
+        ModelSpecs specs = ModelSpecs.of();
         
         Tensor embedding = weights.get("wte.weight"); // embedding  -> [vocab, dim]
         Tensor posEncode = weights.get("wpe.weight"); // pos encode -> [length, dim]
@@ -44,9 +45,9 @@ public Model buildModel(JsonObject config, Map<String, Tensor> weights) {
         vocabLayer.setWeights(embedding.transpose());
         posEncodeLayer.setWeights(posEncode);
         
-        seq.add(new InputLayer(-1));
-        seq.add(embeddingLayer);
-        seq.add(posEncodeLayer);
+        specs.add(new InputLayer(-1).freeze());
+        specs.add(embeddingLayer.freeze());
+        specs.add(posEncodeLayer.freeze());
         
         for (int i = 0; i < layers; i++) {
             String prefix = String.format("h.%s.", i);
@@ -91,7 +92,7 @@ public Model buildModel(JsonObject config, Map<String, Tensor> weights) {
             attention.setOutProj(attnOutWeight);
             attention.setOutBias(attnOutBias);
             
-            seq.add(decoder);
+            specs.add(decoder.freeze());
         }
         
         TokenSelectionLayer selectionLayer = new TokenSelectionLayer();
@@ -103,11 +104,11 @@ public Model buildModel(JsonObject config, Map<String, Tensor> weights) {
         normLayer.setWeights(lnGamma);
         normLayer.setBias(lnBeta);
         
-        seq.add(normLayer);
-        seq.add(selectionLayer);
-        seq.add(vocabLayer);
+        specs.add(normLayer.freeze());
+        specs.add(selectionLayer.freeze());
+        specs.add(vocabLayer.freeze());
         
-        return seq;
+        return specs.compile();
     }
     
     static class TokenSelectionLayer extends Layer {
@@ -118,7 +119,7 @@ public Tensor[] forward(StatesCache cache, Tensor... inputs) {
             Tensor input = inputs[0]; // [batch, seq_len, dim]
             int seqLength = input.shapeAt(1);
             
-            return new Tensor[] { input.slice(Range.all(), Range.point(seqLength - 1), Range.all()).squeeze(1) };
+            return new Tensor[] { input.slice(Range.all(), Range.point(seqLength - 1), Range.all()).squeezeGrad(1) };
         }
         
         @Override
diff --git a/brain4j-llm/src/main/java/org/brain4j/llm/core/loader/ModelLoader.java b/brain4j-llm/src/main/java/org/brain4j/llm/core/loader/ModelLoader.java
@@ -11,7 +11,6 @@
 import org.brain4j.llm.core.model.LLM;
 import org.brain4j.llm.download.callback.ProgressCallback;
 import org.brain4j.llm.download.manager.DownloadManager;
-import org.brain4j.math.commons.result.Result;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -56,7 +55,7 @@ public Tokenizer loadTokenizer(String tokenizerId) throws Exception {
     public Tokenizer loadTokenizer(String tokenizerId, LoadConfig config) throws Exception {
         log.info("Loading tokenizer: {}", tokenizerId);
         
-        ModelInfo info = client.getModelInfo(tokenizerId).unwrap();
+        ModelInfo info = client.getModelInfo(tokenizerId);
         log.debug("Tokenizer info retrieved for: {} (resolved id: {})", tokenizerId, info.id());
         
         String fileToDownload = "tokenizer.json";
@@ -75,7 +74,7 @@ public Tokenizer loadTokenizer(String tokenizerId, LoadConfig config) throws Exc
     public LLM loadModel(String modelId, LoadConfig config) throws Exception {
         log.info("Loading model: {}", modelId);
 
-        ModelInfo info = client.getModelInfo(modelId).unwrap();
+        ModelInfo info = client.getModelInfo(modelId);
         log.debug("Model info retrieved for: {} (resolved id: {})", modelId, info.id());
 
         List<String> filesToDownload = determineFilesToDownload(info, config);
diff --git a/brain4j-llm/src/main/java/org/brain4j/llm/core/model/InferenceProvider.java b/brain4j-llm/src/main/java/org/brain4j/llm/core/model/InferenceProvider.java
@@ -5,6 +5,7 @@
 import java.util.function.Consumer;
 
 public interface InferenceProvider {
+    void move(Device device);
     String chat(String prompt);
     String chat(String prompt, SamplingConfig config);
     String chat(String prompt, SamplingConfig config, Consumer<String> tokenConsumer);
diff --git a/brain4j-llm/src/main/java/org/brain4j/llm/core/model/LLM.java b/brain4j-llm/src/main/java/org/brain4j/llm/core/model/LLM.java
@@ -68,6 +68,11 @@ public LLM compile() throws IOException {
         return this;
     }
     
+    @Override
+    public void move(Device device) {
+        this.model = model.fork(device);
+    }
+    
     @Override
     public String chat(String prompt) {
         return chat(prompt, SamplingConfig.defaultConfig());
diff --git a/brain4j-llm/src/test/java/TestModel.java b/brain4j-llm/src/test/java/TestModel.java
@@ -1,36 +1,62 @@
+import org.brain4j.core.Brain4J;
 import org.brain4j.llm.Models;
 import org.brain4j.llm.core.model.LLM;
 import org.brain4j.llm.core.model.SamplingConfig;
+import org.brain4j.math.gpu.device.Device;
 
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
 
 public class TestModel {
 
     public static void main(String[] args) throws Exception {
         LLM llm = Models.loadModel("gpt2");
-        llm.model().summary();
-        String prompt = "Hello, my name is";
-
-        AtomicLong lastToken = new AtomicLong(System.nanoTime());
-        AtomicReference<Double> total = new AtomicReference<>(0.0);
-        AtomicInteger generated = new AtomicInteger(0);
 
-        System.out.print(prompt);
         SamplingConfig config = SamplingConfig.builder().maxLength(256).build();
-        String response = llm.chat(prompt, config, token -> {
-            long now = System.nanoTime();
-            double took = (now - lastToken.get()) / 1e6;
-            System.out.print(token);
-
-            total.updateAndGet(v -> v + took);
-            lastToken.set(now);
-            generated.incrementAndGet();
-        });
-        double average = total.get() / generated.get();
+        TokenHandler handler = new TokenHandler();
+        String prompt = "Hello, my name is";
 
-        System.out.println("Total ms spent generating = " + total.get());
-        System.out.println("Average ms/token = " + average);
+        Device device = Brain4J.firstDevice();
+        if (device != null) {
+            System.out.printf("Using device %s %n", device.name());
+            llm.move(device);
+        }
+        
+        llm.getModel().summary();
+        llm.chat(prompt, config, handler);
+        handler.printStats();
+    }
+    
+    private static class TokenHandler implements Consumer<String> {
+        
+        private long lastTokenTime;
+        private double totalTime;
+        private int generatedTokens;
+        
+        public TokenHandler() {
+            this.lastTokenTime = System.nanoTime();
+        }
+        
+        @Override
+        public void accept(String s) {
+            long now = System.nanoTime();
+            double took = (now - lastTokenTime) / 1e6;
+            System.out.print(s);
+            
+            this.lastTokenTime = now;
+            this.totalTime += took;
+            this.generatedTokens++;
+        }
+        
+        public void printStats() {
+            double average = totalTime / generatedTokens;
+            
+            System.out.println();
+            System.out.printf("%s generated tokens %n", generatedTokens);
+            System.out.printf("total ms  = %.2f %n", totalTime);
+            System.out.printf("avg/token = %.2f %n", average);
+        }
     }
 }
diff --git a/brain4j-math/src/main/java/org/brain4j/math/activation/Activation.java b/brain4j-math/src/main/java/org/brain4j/math/activation/Activation.java
@@ -72,7 +72,7 @@ default Tensor activate(Tensor input) {
             GpuTensor result = new GpuTensor(device, gpuInput.shape());
 
             try (GpuQueue queue = GpuContext.getOrCreate(device)) {
-                long kernel = GpuContext.kernel(device, getKernelPrefix() + "_forward");
+                long kernel = GpuContext.findKernel(device, getKernelPrefix() + "_forward");
 
                 KernelFactory factory = createKernel(kernel, gpuInput, result);
                 factory.launch(queue, 1, gpuInput.size());
@@ -108,7 +108,7 @@ default Tensor derivative(Tensor input) {
             GpuTensor result = new GpuTensor(device, gpuInput.shape());
 
             try (GpuQueue queue = GpuContext.getOrCreate(device)) {
-                long kernel = GpuContext.kernel(device, getKernelPrefix() + "_backward");
+                long kernel = GpuContext.findKernel(device, getKernelPrefix() + "_backward");
 
                 KernelFactory factory = createKernel(kernel, gpuInput, result);
                 factory.launch(queue, 1, gpuInput.size());
diff --git a/brain4j-math/src/main/java/org/brain4j/math/clipper/impl/HardClipper.java b/brain4j-math/src/main/java/org/brain4j/math/clipper/impl/HardClipper.java
@@ -26,7 +26,7 @@ public void clipCpu(CpuTensor grad) {
     @Override
     public void clipGpu(GpuTensor grad) {
         Device device = grad.device();
-        long kernel = GpuContext.kernel(device, kernelName());
+        long kernel = GpuContext.findKernel(device, kernelName());
         
         try (GpuQueue queue = GpuContext.getOrCreate(device)) {
             KernelFactory.create(kernel)
diff --git a/brain4j-math/src/main/java/org/brain4j/math/gpu/GpuContext.java b/brain4j-math/src/main/java/org/brain4j/math/gpu/GpuContext.java
@@ -28,7 +28,7 @@ public static void register(Device device, String kernelName, long program) {
             });
     }
 
-    public static long kernel(Device device, String kernelName) {
+    public static long findKernel(Device device, String kernelName) {
         Map<String, Long> deviceKernels = kernelCache.get(device);
 
         if (deviceKernels == null) {
diff --git a/brain4j-math/src/main/java/org/brain4j/math/gpu/kernel/KernelFactory.java b/brain4j-math/src/main/java/org/brain4j/math/gpu/kernel/KernelFactory.java
@@ -26,7 +26,7 @@ public static KernelFactory create(long kernel) {
     }
 
     public static KernelFactory create(Device device, String kernelName) {
-        return create(GpuContext.kernel(device, kernelName));
+        return create(GpuContext.findKernel(device, kernelName));
     }
 
     public KernelFactory addIntParam(int variable) {
diff --git a/brain4j-math/src/main/java/org/brain4j/math/tensor/Tensor.java b/brain4j-math/src/main/java/org/brain4j/math/tensor/Tensor.java
@@ -870,4 +870,8 @@ default Tensor sigmoid() {
     default Tensor tanh() {
         return activateAuto(new TanhActivation());
     }
+    
+    default Tensor cpu() {
+        return to(null);
+    }
 }
diff --git a/brain4j-math/src/main/java/org/brain4j/math/tensor/impl/BaseTensor.java b/brain4j-math/src/main/java/org/brain4j/math/tensor/impl/BaseTensor.java
@@ -467,10 +467,7 @@ public Tensor squeeze() {
             }
         }
         
-        Tensor reshaped = reshape(newShape);
-        reshaped.setAutogradContext(autogradContext);
-        
-        return reshaped;
+        return reshape(newShape);
     }
 
     @Override
@@ -493,11 +490,8 @@ public Tensor squeeze(int dim) {
                 newShape[idx++] = shape[i];
             }
         }
-
-        Tensor result = reshape(newShape);
-        result.setAutogradContext(autogradContext);
         
-        return result;
+        return reshape(newShape);
     }
 
     @Override

Original file line number	Diff line number	Diff line change
`@@ -123,7 +123,11 @@ public static Device firstDevice() {`
`123`	`123`	`throw new IllegalStateException("No GPU-acceleration device has been found!");`
`124`	`124`	`}`
`125`	`125`
`126`		`- return DeviceUtils.findDevice(devices.getFirst());`
	`126`	`+ Device device = DeviceUtils.findDevice(devices.getFirst());`
	`127`	`+`
	`128`	`+ if (device != null) Brain4J.initKernels(device);`
	`129`	`+`
	`130`	`+ return device;`
`127`	`131`	`}`
`128`	`132`
`129`	`133`	`/**`
Original file line number	Diff line number	Diff line change
`@@ -59,7 +59,8 @@ public Tensor[] predict(StatesCache cache, Tensor... inputs) {`
`59`	`59`	`input = input.reshape(1, input.elements()); // reshape to [batch, input_size]`
`60`	`60`	`}`
`61`	`61`
`62`		`- buffer[i] = cache.isTraining() ? input.withGrad() : input;`
	`62`	`+ Tensor chosen = cache.isTraining() ? input.withGrad() : input;`
	`63`	`+ buffer[i] = chosen.to(device);`
`63`	`64`	`}`
`64`	`65`
`65`	`66`	`for (Layer layer : layers) {`
Original file line number	Diff line number	Diff line change
`@@ -114,7 +114,7 @@ public Tensor[] forward(StatesCache cache, Tensor... inputs) {`
`114`	`114`	`}`
`115`	`115`
`116`	`116`	`Range[] slicingRanges = {`
`117`		`- Range.all(), Range.point(seqLength - 1), Range.all()`
	`117`	`+ Range.all(), Range.point(seqLength - 1), Range.all()`
`118`	`118`	`}; // [batch, 1, dim]`
`119`	`119`	`Tensor cachedOutput = cache.get(outProj);`
`120`	`120`	`Tensor cachedQKV = cache.get(weights);`
Original file line number	Diff line number	Diff line change
`@@ -66,11 +66,8 @@ public Tensor attend(Tensor input) {`
`66`	`66`	`}`
`67`	`67`
`68`	`68`	`if (context != null) {`
`69`		`- return training`
`70`		`- ? context.squeezeGrad(1) // [B,L,d]`
`71`		`- : context.squeeze(1);`
	`69`	`+ return context.squeezeGrad(1); // [B,L,d]`
`72`	`70`	`}`
`73`		`- // fallthrough if context null`
`74`	`71`	`}`
`75`	`72`
`76`	`73`	`// fallback to standard path with autograd support`