The update code could be copy pasted (almost) 
It seems like there is some boiler plate code that is missing in the trax example to help with creating the buffers (not sure where that code is).
But then that’s it, if you remove the input params validation and empty lines between the math.
We are actually refactoring the optimizers right now to move the update code outside of the state tracking. As you can see in the PR here 21 the update function is not longer than this one.