Inline PassthroughVariable

This Variable should not be used outside the ComputationContext

Author: FlorentinD

ml/ml-core/src/main/java/org/neo4j/gds/ml/core/ComputationContext.java

@@ -20,7 +20,7 @@
 package org.neo4j.gds.ml.core;
 
 import org.jetbrains.annotations.TestOnly;
-import org.neo4j.gds.ml.core.functions.PassthroughVariable;
+import org.neo4j.gds.ml.core.functions.SingleParentVariable;
 import org.neo4j.gds.ml.core.tensor.Tensor;
 import org.neo4j.gds.ml.core.tensor.TensorFactory;
 
@@ -71,7 +71,7 @@ public void backward(Variable<?> function) {
 
         gradients.clear();
         Queue<BackPropTask> executionQueue = new LinkedBlockingQueue<>();
-        PassthroughVariable<?> dummy = new PassthroughVariable<>(function);
+        var dummy = new PassthroughVariable<>(function);
         executionQueue.add(new BackPropTask(function, dummy));
         Map<Variable<?>, AtomicInteger> upstreamCounters = new HashMap<>();
         initUpstream(dummy, upstreamCounters);
@@ -167,4 +167,25 @@ static class BackPropTask {
         }
     }
 
+    private static class PassthroughVariable<T extends Tensor<T>> extends SingleParentVariable<T, T> {
+
+        public PassthroughVariable(Variable<T> parent) {
+            super(parent, parent.dimensions());
+
+            if (parent instanceof PassthroughVariable) {
+                throw new IllegalArgumentException("Redundant use of PassthroughVariables. Chaining does not make sense.");
+            }
+        }
+
+        @Override
+        public T apply(ComputationContext ctx) {
+            return ctx.data(parent);
+        }
+
+        @Override
+        public T gradientForParent(ComputationContext ctx) {
+            // initialize gradient computation with `1`
+            return ctx.data(parent).map(v -> 1);
+        }
+    }
 }