### Question:
AdamOptimizer { public synchronized void update(ComputationContext otherCtx) { iteration += 1; variables.forEach(variable -> otherCtx.gradient(variable).mapInPlace(this::clip)); momentumTerms = IntStream.range(0, variables.size()) .mapToObj(i -> { Variable<?> variable = variables.get(i); Tensor<?> momentumTerm = momentumTerms.get(i); return castAndAdd( momentumTerm.scalarMultiply(beta_1), otherCtx.gradient(variable).scalarMultiply(1 - beta_1) ); }) .collect(Collectors.toList()); velocityTerms = IntStream.range(0, variables.size()) .mapToObj(i -> { Variable<?> variable = variables.get(i); Tensor<?> velocityTerm = velocityTerms.get(i); Tensor<?> gradient = otherCtx.gradient(variable); Tensor<?> squaredGradient = gradient.elementwiseProduct(gradient); return castAndAdd( velocityTerm.scalarMultiply(beta_2), squaredGradient.scalarMultiply(1 - beta_2) ); }) .collect(Collectors.toList()); List<Tensor<?>> mCaps = momentumTerms.stream() .map(mTerm -> mTerm.scalarMultiply(1d / (1 - Math.pow(beta_1, iteration)))) .collect(Collectors.toList()); List<Tensor<?>> vCaps = velocityTerms.stream() .map(vTerm -> vTerm.scalarMultiply(1d / (1 - Math.pow(beta_2, iteration)))) .collect(Collectors.toList()); IntStream.range(0, variables.size()) .forEach(i -> { Weights<? extends Tensor<?>> variable = variables.get(i); Tensor<?> mCap = mCaps.get(i); Tensor<?> vCap = vCaps.get(i); Tensor<?> theta_0 = variable.data(); theta_0.addInPlace(mCap .scalarMultiply(-alpha) .elementwiseProduct(vCap.map(v -> 1 / (Math.sqrt(v) + epsilon))) ); }); } AdamOptimizer(List<Weights<? extends Tensor<?>>> variables);  AdamOptimizer(
        List<Weights<? extends Tensor<?>>> variables,
        double learningRate
    ); synchronized void update(ComputationContext otherCtx);  }### Answer:
@Test void shouldConverge() { Weights<Matrix> weights = new Weights<>(new Matrix( new double[]{ 0.1, 0.1, 0.1, 0.4, 0.3, 0.9, 0.01, 0.6, 0.5 }, 3, 3 )); Variable<Matrix> expectedOptimum = new Weights<>(new Matrix( new double[]{ 0.11, 0.13, 0.231, 0.4, 0.3, 0.9, 0.01, 0.6, 0.15 }, 3, 3 )); ComputationContext ctx = new ComputationContext(); ctx.forward(weights); AdamOptimizer adam = new AdamOptimizer(List.of(weights)); double oldLoss = Double.MAX_VALUE; while(true) { Variable<Matrix> difference = new MatrixSum(List.of( weights, new ConstantScale<>(expectedOptimum, -1) )); L2Norm lossFunction = new L2Norm(difference); double newLoss = ctx.forward(lossFunction).dataAt(0); double d = oldLoss - newLoss; if (Math.abs(d) < 1e-8) break; oldLoss = newLoss; ctx.backward(lossFunction); adam.update(ctx); } assertTrue(oldLoss < 1e-4, "oldLoss was : " + oldLoss); }