This code is the solution to an edabit problem.
Go see for yourself. It works.
I was watching a video (go watch them, it's pure entertainment) where two guys worked on this problem during a live stream and their solution was unnecessarily complicated. That inspired me to go WAY too far on this and use the Adam optimizer I hacked together on this problem.
Getting this to work within the time constraints set by edabit was a little bit of a challenge, but through some optimization, I got it to converge in under 12 seconds!