make differentiator more robust by pulling out the partials you need after calculation instead of while its visiting vertices
@@ -34,16 +34,10 @@ public class Differentiator {
Map<Vertex, PartialDerivatives> dwrtOf = new HashMap<>();
collectPartials(singletonMap(ofVertex, dWrtOfVertex), dwrtOf, dOfShape);
- Map<VertexId, PartialDerivatives> wrtOf = new HashMap<>();
Vertex<?> visiting;
while ((visiting = priorityQueue.poll()) != null) {
- if (wrt.contains(visiting)) {
- wrtOf.put(visiting.getId(), dwrtOf.get(visiting));
- continue;
- }
-
if (visiting instanceof Differentiable) {
Differentiable visitingDifferentiable = ((Differentiable) visiting);
Map<Vertex, PartialDerivatives> partialDerivatives = visitingDifferentiable.reverseModeAutoDifferentiation(dwrtOf.get(visiting));
@@ -60,6 +54,13 @@ public class Differentiator {
}
}
+ Map<VertexId, PartialDerivatives> wrtOf = new HashMap<>();
+ for (Vertex vertex : wrt) {
+ if (dwrtOf.containsKey(vertex)) {
+ wrtOf.put(vertex.getId(), dwrtOf.get(vertex));
+ }
+ }
+
return wrtOfToOfWrt(wrtOf).get(ofVertex.getId());
}
