I ended up going through the notebook and making a bunch of changes: both to the prose and also to the code. Switched from using gradient descent to find fixed points to using Newton's method which is at least 100x faster. Also, figured out a concise way of expressing the q-map, c-map, and chi at once.
I ended up going through the notebook and making a bunch of changes: both to the prose and also to the code. Switched from using gradient descent to find fixed points to using Newton's method which is at least 100x faster. Also, figured out a concise way of expressing the q-map, c-map, and chi at once.

PiperOrigin-RevId: 291762995

google/neural-tangents