OMG OMG OMG!!!!!! ZOOOOOMIESSS <3333

* IT WORKS FOR GPU NOW, DON'T TOUCH ANYTHING YET LMAO
* Bugs noted: It throws if you don't provide a loss function; should be an extremely trivial fix, but I had this working, broke it, panicked, and had to fix it over the course of half an hour. So please, understand that I HAD TO COMMIT!!!

src/neural-network-gpu.ts

@@ -22,6 +22,16 @@ import {
   NeuralNetwork,
 } from './neural-network';
 import { release } from './utilities/kernel';
+import { LossFunction, LossFunctionInputs, LossFunctionState } from './utilities/loss';
+
+function loss(
+  actual: number,
+  expected: number,
+  inputs: LossFunctionInputs,
+  state: LossFunctionState
+) {
+  return expected - actual;
+}
 
 export interface INeuralNetworkGPUDatumFormatted {
   input: KernelOutput;
@@ -96,8 +106,8 @@ function weightedSumTanh(
   return Math.tanh(sum);
 }
 
-function loss(output: number, target: number): number {
-  return target - output;
+function calcErrorOutput(value: number): number {
+  return value;
 }
 
 function calcDeltasSigmoid(error: number, output: number): number {
@@ -266,6 +276,16 @@ export class NeuralNetworkGPU<
     this.gpu = new GPU({ mode: options.mode });
   }
 
+  get lossFunction(): LossFunction {
+    return typeof this._lossFunction === "function" ? this._lossFunction : loss
+  }
+
+  set lossFunction(
+    value: LossFunction
+  ) {
+    this._lossFunction = value;
+  }
+
   initialize(): void {
     super.initialize();
     this.buildRunInput();
@@ -380,7 +400,7 @@ export class NeuralNetworkGPU<
   };
 
   buildCalculateDeltas(): void {
-    let calcDeltas: GPUFunction<[number, number]>;
+    let calcDeltas: GPUFunction<[number, number, LossFunctionInputs, LossFunctionState]>;
     switch (this.trainOpts.activation) {
       case 'sigmoid':
         calcDeltas = calcDeltasSigmoid;
@@ -400,31 +420,34 @@ export class NeuralNetworkGPU<
         );
     }
 
-    let _loss = typeof this.loss === "function" ? this.loss : loss;
+    const loss: LossFunction = this._lossFunction ?? function loss(actual, expected) { return expected - actual; };
 
     calcDeltas = alias(
       utils.getMinifySafeName(() => calcDeltas),
       calcDeltas
     );
     this.gpu.addFunction(calcDeltas);
+    this.gpu.addFunction(loss);
     for (let layer = this.outputLayer; layer > 0; layer--) {
       if (layer === this.outputLayer) {
         // eslint-disable-next-line @typescript-eslint/ban-ts-comment
         // @ts-expect-error
         this.backwardPropagate[this.outputLayer] = this.gpu.createKernelMap(
           {
-            error: _loss,
+            error: calcErrorOutput,
           },
           function (
             this: IKernelFunctionThis,
             outputs: number[],
-            targets: number[]
+            targets: number[],
+            inputs: LossFunctionInputs,
+            state: LossFunctionState
           ): number {
             const output = outputs[this.thread.x];
             const target = targets[this.thread.x];
             // eslint-disable-next-line @typescript-eslint/ban-ts-comment
             // @ts-expect-error
-            return calcDeltas(loss(output, target), output);
+            return calcDeltas(calcErrorOutput(loss(output, target, inputs, state)), output);
           },
           {
             output: [this.sizes[this.outputLayer]],
@@ -480,7 +503,7 @@ export class NeuralNetworkGPU<
       if (layer === this.outputLayer) {
         // eslint-disable-next-line @typescript-eslint/ban-ts-comment
         // @ts-expect-error
-        output = this.backwardPropagate[layer](this.outputs[layer], target);
+        output = this.backwardPropagate[layer](this.outputs[layer], target, this.outputs[0], this.lossState);
       } else {
         // eslint-disable-next-line @typescript-eslint/ban-ts-comment
         // @ts-expect-error

src/neural-network.ts

@@ -13,7 +13,6 @@ import { mse } from './utilities/mse';
 import { randos } from './utilities/randos';
 import { zeros } from './utilities/zeros';
 import { LossFunction, LossFunctionInputs, LossFunctionState } from './utilities/loss';
-
 type NeuralNetworkFormatter =
   | ((v: INumberHash) => Float32Array)
   | ((v: number[]) => Float32Array);
@@ -41,6 +40,15 @@ export function getTypedArrayFn(
   };
 }
 
+function loss(
+  actual: number,
+  expected: number,
+  inputs: LossFunctionInputs,
+  state: LossFunctionState
+) {
+  return expected - actual;
+}
+
 export type NeuralNetworkActivation =
   | 'sigmoid'
   | 'relu'
@@ -110,7 +118,7 @@ export interface INeuralNetworkTrainOptions {
   errorThresh: number;
   log: boolean | ((status: INeuralNetworkState) => void);
   logPeriod: number;
-  loss: boolean | LossFunction;
+  loss?: LossFunction;
   lossState?: LossFunctionState;
   lossStateSize: number;
   leakyReluAlpha: number;
@@ -132,7 +140,7 @@ export function trainDefaults(): INeuralNetworkTrainOptions {
     errorThresh: 0.005, // the acceptable error percentage from training data
     log: false, // true to use console.log, when a function is supplied it is used
     logPeriod: 10, // iterations between logging out
-    loss: false,
+    loss,
     lossStateSize: 1,
     leakyReluAlpha: 0.01,
     learningRate: 0.3, // multiply's against the input and the delta then adds to momentum
@@ -191,15 +199,15 @@ export class NeuralNetwork<
     return this.runInput(input);
   };
 
-  calculateDeltas: (output: Float32Array, input: LossFunctionInputs) => void = (
+  calculateDeltas: (output: Float32Array, input: Float32Array) => void = (
     output: Float32Array,
-    input: LossFunctionInputs
+    input: Float32Array
   ): void => {
     this.setActivation();
     return this.calculateDeltas(output, input);
   };
 
-  loss: boolean | LossFunction = false;
+  _lossFunction?: LossFunction;
 
   // adam
   biasChangesLow: Float32Array[] = [];
@@ -690,7 +698,7 @@ export class NeuralNetwork<
     data: Array<INeuralNetworkDatum<Partial<InputType>, Partial<OutputType>>>,
     options: Partial<INeuralNetworkTrainOptions> = {}
   ): INeuralNetworkState {
-    this.loss = options.loss ?? false;
+    this._lossFunction ??= options.loss;
 
     const { preparedData, status, endTime } = this.prepTraining(
       data as Array<INeuralNetworkDatum<InputType, OutputType>>,
@@ -747,7 +755,7 @@ export class NeuralNetwork<
     return null;
   }
 
-  _calculateDeltasSigmoid(target: Float32Array, input: LossFunctionInputs): void {
+  _calculateDeltasSigmoid(target: Float32Array, input: Float32Array): void {
     for (let layer = this.outputLayer; layer >= 0; layer--) {
       const activeSize = this.sizes[layer];
       const activeOutput = this.outputs[layer];
@@ -760,7 +768,7 @@ export class NeuralNetwork<
 
         let error = 0;
         if (layer === this.outputLayer) {
-          if (typeof this.loss === "function") error = this.loss.call({thread: {x: node}}, output, target[node], input, this.lossState);
+          if (typeof this._lossFunction === "function") error = this._lossFunction.call({thread: {x: node}}, output, target[node], input, this.lossState);
           else error = target[node] - output;
         } else {
           const deltas = this.deltas[layer + 1];
@@ -774,7 +782,7 @@ export class NeuralNetwork<
     }
   }
 
-  _calculateDeltasRelu(target: Float32Array, input: LossFunctionInputs): void {
+  _calculateDeltasRelu(target: Float32Array, input: Float32Array): void {
     for (let layer = this.outputLayer; layer >= 0; layer--) {
       const currentSize = this.sizes[layer];
       const currentOutputs = this.outputs[layer];
@@ -788,7 +796,7 @@ export class NeuralNetwork<
 
         let error = 0;
         if (layer === this.outputLayer) {
-          if (typeof this.loss === "function") error = this.loss.call({thread: {x: node}}, output, target[node], input, this.lossState);
+          if (typeof this._lossFunction === "function") error = this._lossFunction.call({thread: {x: node}}, output, target[node], input, this.lossState);
           else error = target[node] - output;
         } else {
           for (let k = 0; k < nextDeltas.length; k++) {
@@ -801,7 +809,7 @@ export class NeuralNetwork<
     }
   }
 
-  _calculateDeltasLeakyRelu(target: Float32Array, input: LossFunctionInputs): void {
+  _calculateDeltasLeakyRelu(target: Float32Array, input: Float32Array): void {
     const alpha = this.trainOpts.leakyReluAlpha;
     for (let layer = this.outputLayer; layer >= 0; layer--) {
       const currentSize = this.sizes[layer];
@@ -816,7 +824,7 @@ export class NeuralNetwork<
 
         let error = 0;
         if (layer === this.outputLayer) {
-          if (typeof this.loss === "function") error = this.loss.call({thread: {x: node}}, output, target[node], input, this.lossState);
+          if (typeof this._lossFunction === "function") error = this._lossFunction.call({thread: {x: node}}, output, target[node], input, this.lossState);
           else error = target[node] - output;
         } else {
           for (let k = 0; k < nextDeltas.length; k++) {
@@ -829,7 +837,7 @@ export class NeuralNetwork<
     }
   }
 
-  _calculateDeltasTanh(target: Float32Array, input: LossFunctionInputs): void {
+  _calculateDeltasTanh(target: Float32Array, input: Float32Array): void {
     for (let layer = this.outputLayer; layer >= 0; layer--) {
       const currentSize = this.sizes[layer];
       const currentOutputs = this.outputs[layer];
@@ -843,7 +851,7 @@ export class NeuralNetwork<
 
         let error = 0;
         if (layer === this.outputLayer) {
-          if (typeof this.loss === "function") error = this.loss.call({thread: {x: node}}, output, target[node], input, this.lossState);
+          if (typeof this._lossFunction === "function") error = this._lossFunction.call({thread: {x: node}}, output, target[node], input, this.lossState);
           else error = target[node] - output;
         } else {
           for (let k = 0; k < nextDeltas.length; k++) {

src/utilities/loss.ts

@@ -1,4 +1,4 @@
-export type LossFunctionInputs = number[] | Float32Array;
+export type LossFunctionInputs = number[] | number[][] | number[][][] | Float32Array | Float32Array[] | Float32Array[][];
 
 export type LossFunctionState = number[][][] | Float32Array[][];
 

src/utilities/to-svg.ts

@@ -467,6 +467,7 @@ export function toSVG<
   // Get network size array for NeuralNetwork or NeuralNetworkGPU
   let sizes: number[] = [];
   if (net instanceof NeuralNetwork || net instanceof NeuralNetworkGPU) {
+    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
     // @ts-ignore
     sizes = getNeuralNetworkSizes(net);
   }