Consider self gradient in CrossEntropyLoss::gradient

Author: FlorentinD

ml/ml-core/src/main/java/org/neo4j/gds/ml/core/functions/CrossEntropyLoss.java

@@ -71,16 +71,14 @@ public Tensor<?> gradient(Variable<?> parent, ComputationContext ctx) {
             Matrix gradient = predictionsMatrix.createWithSameDimensions();
             var targetsVector = ctx.data(targets);
 
-            var multiplier = -1.0 / gradient.rows();
+            var multiplier = - ctx.gradient(this).value() / gradient.rows();
             for (int row = 0; row < gradient.rows(); row++) {
                 var trueClass = (int) targetsVector.dataAt(row);
                 var predictedProbabilityForTrueClass = predictionsMatrix.dataAt(row * predictionsMatrix.cols() + trueClass);
 
                 // Compare to a threshold value rather than `0`, very small probability can result in setting infinite gradient values.
                 if (predictedProbabilityForTrueClass > PREDICTED_PROBABILITY_THRESHOLD) {
-                    gradient.setDataAt(
-                        row * predictionsMatrix.cols() + trueClass,
-                        multiplier / predictedProbabilityForTrueClass
+                    gradient.setDataAt(row, trueClass, multiplier / predictedProbabilityForTrueClass
                     );
                 }
             }

ml/ml-core/src/test/java/org/neo4j/gds/ml/core/functions/CrossEntropyLossTest.java

@@ -76,8 +76,28 @@ void shouldComputeGradientCorrectly() {
     }
 
     @Test
-    void infiniteSmallProbabilities() {
+    void considerSelfGradient() {
+        var targets = Constant.vector(new double[]{1.0, 2.0, 0.0});
+        var predictions = new Weights<>(
+            new Matrix(
+                new double[]{
+                    0.35, 0.65, 0.0,
+                    0.45, 0.45, 0.1,
+                    0.14, 0.66, 0.2
+                },
+                3, 3
+            )
+        );
 
+        var loss = new CrossEntropyLoss(predictions, targets);
+        var chainedLoss = new Sigmoid<>(loss);
+
+        finiteDifferenceShouldApproximateGradient(predictions, chainedLoss);
+    }
+
+
+    @Test
+    void infiniteSmallProbabilities() {
         var predictions = new Weights<>(new Matrix(new double[]{5.277E-321, 5.277E-321}, 1, 2));
         var targets = Constant.vector(new double[]{1});