general code refactors

Author: xEcho1337

brain4j-core/src/main/java/org/brain4j/core/layer/Layer.java

@@ -106,18 +106,27 @@ public void deserialize(DataInputStream stream) throws Exception {
         }
     }
 
-    public void compile(WeightInitializer weightInit, LossFunction lossFunction, Optimizer optimizer, Updater updater) {
+    public void compile(
+        WeightInitializer weightInit,
+        Optimizer optimizer,
+        Updater updater
+    ) {
         this.weightInit = weightInit;
-        this.lossFunction = lossFunction;
         this.optimizer = optimizer;
         this.updater = updater;
     }
 
-    public Tensor computeLoss(StatesCache cache, Tensor targets, Tensor outputs, LossFunction lossFunction) {
+    public Tensor computeLoss(
+        int index,
+        StatesCache cache,
+        Tensor targets,
+        Tensor outputs,
+        LossFunction lossFunction
+    ) {
         Tensor error = outputs.minus(targets);
         Tensor derivatives = activation.getDerivative(outputs);
 
-        Tensor input = cache.getInputTensor(this);
+        Tensor input = cache.getInputTensor(index);
         Tensor delta = lossFunction.getDelta(error, derivatives);
 
         Tensor weightsGradient = input.transpose().matmul(delta);
@@ -128,7 +137,11 @@ public Tensor computeLoss(StatesCache cache, Tensor targets, Tensor outputs, Los
         return delta;
     }
 
-    public void connect(Random generator, Layer previous, double bound) {
+    public void connect(
+        Random generator,
+        Layer previous,
+        double bound
+    ) {
         if (previous == null) return;
 
         int input = previous.getTotalNeurons();
@@ -145,9 +158,14 @@ public void connect(Random generator, Layer previous, double bound) {
         }
     }
 
-    public abstract Tensor forward(StatesCache cache, Tensor input, boolean training);
+    public abstract Tensor forward(
+        int index,
+        StatesCache cache,
+        Tensor input,
+        boolean training
+    );
 
-    public Tensor backward(StatesCache cache, Layer previous, Tensor delta) {
+    public Tensor backward(int index, StatesCache cache, Layer previous, Tensor delta) {
         throw new UnsupportedOperationException("Not implemented for " + this.getClass().getSimpleName());
     }
 

brain4j-core/src/main/java/org/brain4j/core/layer/impl/BatchNorm.java

@@ -52,7 +52,12 @@ public boolean canPropagate() {
     }
 
     @Override
-    public Tensor forward(StatesCache cache, Tensor input, boolean training) {
+    public Tensor forward(
+        int index,
+        StatesCache cache,
+        Tensor input,
+        boolean training
+    ) {
         int batchSize = input.shape()[0];
 
         Tensor transposed = input.transpose(); // [dimension, batch_size]

brain4j-core/src/main/java/org/brain4j/core/layer/impl/DenseLayer.java

@@ -38,7 +38,12 @@ public String getLayerName() {
     }
 
     @Override
-    public Tensor forward(StatesCache cache, Tensor input, boolean training) {
+    public Tensor forward(
+        int index,
+        StatesCache cache,
+        Tensor input,
+        boolean training
+    ) {
         Tensor Z = input.matmul(weights); // [batch_size, n_out]
 
         int batchSize = Z.shape()[0];
@@ -54,7 +59,7 @@ public Tensor forward(StatesCache cache, Tensor input, boolean training) {
         }
 
         if (nextLayer instanceof LayerNorm layerNorm) {
-            Z = layerNorm.forward(cache, Z, training);
+            Z = layerNorm.forward(index, cache, Z, training);
         }
 
         Tensor activated = activation.activate(Z);
@@ -64,13 +69,18 @@ public Tensor forward(StatesCache cache, Tensor input, boolean training) {
     }
 
     @Override
-    public Tensor backward(StatesCache cache, Layer previous, Tensor delta) {
-        Tensor input = cache.getInputTensor(this);
-        Tensor output = cache.getOutputTensor(this);
+    public Tensor backward(
+        int index,
+        StatesCache cache,
+        Layer previous,
+        Tensor delta
+    ) {
+        Tensor input = cache.getInputTensor(index);
+        Tensor output = cache.getOutputTensor(index);
         Tensor derivative = activation.getDerivative(output); // [batch_size, n_out]
 
-        Tensor weightsNext = previous.getWeights();  // [n_out, n_out_next]
-        Tensor deltaProjected = delta.matmul(weightsNext.transpose());  // [batch_size x n_out]
+        Tensor weightsNext = previous.getWeights(); // [n_out, n_out_next]
+        Tensor deltaProjected = delta.matmul(weightsNext.transpose()); // [batch_size x n_out]
 
         Tensor deltaThisLayer = deltaProjected.mul(derivative); // [batch_size x n_out]
 

brain4j-core/src/main/java/org/brain4j/core/layer/impl/DropoutLayer.java

@@ -67,12 +67,18 @@ public boolean canPropagate() {
      * during training. Meanwhile it will scale the input tensor by {@code 1 - dropout} during inferencing.
      *
      * @param nextLayer
+     * @param index
      * @param input     The input tensor.
      * @param training  If it's called during training.
      * @return The resulting tensor.
      */
     @Override
-    public Tensor forward(StatesCache cache, Tensor input, boolean training) {
+    public Tensor forward(
+        int index,
+        StatesCache cache,
+        Tensor input,
+        boolean training
+    ) {
         if (training) {
             return scale(input);
         }

brain4j-core/src/main/java/org/brain4j/core/layer/impl/LayerNorm.java

@@ -57,7 +57,12 @@ public boolean canPropagate() {
     }
 
     @Override
-    public Tensor forward(StatesCache cache, Tensor input, boolean training) {
+    public Tensor forward(
+        int index,
+        StatesCache cache,
+        Tensor input,
+        boolean training
+    ) {
         int batchSize = input.shape()[0];
         Tensor result = input.clone();
 

brain4j-core/src/main/java/org/brain4j/core/layer/impl/conv/ConvLayer.java

@@ -20,7 +20,8 @@ public class ConvLayer extends Layer {
     private int padding;
     private int stride;
 
-    public ConvLayer(Activations activation, int filters, int filtersWidth, int filtersHeight) {
+    public ConvLayer(Activations activation, int filters, int filtersWidth, int filtersHeight
+    ) {
         this(activation.getFunction(), filters, filtersWidth, filtersHeight, 1, 1);
     }
 
@@ -36,7 +37,11 @@ public ConvLayer(Activation activation, int filters, int filtersWidth, int filte
     }
 
     @Override
-    public void connect(Random generator, Layer previous, double bound) {
+    public void connect(
+        Random generator,
+        Layer previous,
+        double bound
+    ) {
         if (previous instanceof ConvLayer convLayer) {
             this.channels = convLayer.getChannels();
         } else if (previous instanceof InputLayer inputLayer) {
@@ -79,7 +84,7 @@ public void deserialize(DataInputStream stream) throws Exception {
     }
 
     @Override
-    public Tensor forward(StatesCache cache, Tensor input, boolean training) {
+    public Tensor forward(int index, StatesCache cache, Tensor input, boolean training) {
         // [batch_size, channels, height, width]
         return input.convolve(weights)
                 .map(x -> activation.activate(x));

brain4j-core/src/main/java/org/brain4j/core/layer/impl/conv/FlattenLayer.java

@@ -7,12 +7,7 @@
 public class FlattenLayer extends Layer {
 
     @Override
-    public int getTotalNeurons() {
-        return super.getTotalNeurons();
-    }
-
-    @Override
-    public Tensor forward(StatesCache cache, Tensor input, boolean training) {
+    public Tensor forward(int index, StatesCache cache, Tensor input, boolean training) {
         return input.reshape(1, input.elements());
     }
 }

brain4j-core/src/main/java/org/brain4j/core/layer/impl/conv/InputLayer.java

@@ -17,7 +17,7 @@ public InputLayer(int width, int height, int channels) {
     }
 
     @Override
-    public Tensor forward(StatesCache cache, Tensor input, boolean training) {
+    public Tensor forward(int index, StatesCache cache, Tensor input, boolean training) {
         return input;
     }
 

brain4j-core/src/main/java/org/brain4j/core/layer/impl/conv/pooling/AveragePooling.java

@@ -26,7 +26,7 @@ public boolean canConnect() {
     }
 
     @Override
-    public Tensor forward(StatesCache cache, Tensor input, boolean training) {
+    public Tensor forward(int index, StatesCache cache, Tensor input, boolean training) {
         return null;
     }
 

brain4j-core/src/main/java/org/brain4j/core/layer/impl/conv/pooling/MaxPooling.java

@@ -26,7 +26,7 @@ public boolean canConnect() {
     }
 
     @Override
-    public Tensor forward(StatesCache cache, Tensor input, boolean training) {
+    public Tensor forward(int index, StatesCache cache, Tensor input, boolean training) {
         return null;
     }