Adam: Running the Full Optimization Loop

The problem, from first principles

02-adam-single-step computes one Adam update. Training repeats that update many times, and the moments must carry over from one step to the next. This question wraps the single step in the loop that training actually runs, then uses it to minimize a simple function you can check by eye: f(x) = x², whose gradient is 2x and whose minimum is at 0.

From theory to code

Implement adam_minimize(grad_fn, theta0, steps, lr, beta1, beta2, eps). It runs steps Adam updates starting from theta0 and returns the full trajectory: the starting value followed by the value after each update. The signature and docstring are already in the editor.

Constraints

• The trajectory has length steps + 1.
• grad_fn takes the current theta and returns the gradient at that point.
• Works on scalars and NumPy arrays.

Hints