Docs cleaning + JAX implementation of the glm models

[ryan-odea]

• Cleaned up some of the docs which would presumably be automatically inferred from sphinx.autodoc
• Added JAX as a backend glm system (note: only attaches to the GPU on linux systems). Seems to work pretty well. I am still rather naive in my pytree implementation so some speed is probably left on the table
• changed .summary2() to .summary() forwarding to just .summary() - matches a bit more on what a user might want from the API (and matches OLS...summary()). Perhaps a bit more intuitive, but open to discussion

[ryan-odea]

Was clicking around a little too quickly and accidentally merged immediately in from my fork. We can still discuss it here (and revert if needed!) Sorry for the hassle :(

[remlapmot]

This is great - really fascinating.

I am really interested to try this (I have WSL on my Windows machine at home, and access to our Isambard AI cluster, will try to try soon).

I had been looking for something similar in R - I found a GitHub only package which just runs on Apple (Accelerate framework) - https://github.com/hughjonesd/Rmlx/ - do you know of anything similar in R?

For simplicity I will guard the failing tests (i.e., currently all the tests in CI are on on macos - maybe I'll add some on Linux) on my PR branch, so I'll close this.