Type checking, testing, code formatting

[matomatical]

This library is getting bigger than the typical couple-of-files scale I usually work at, and involves working with complex structures and multi-dimensional arrays. So far I've been able to stay in control, with the help of JAX transformations like vmap and careful documentation including leaning into Python type annotations more than I have before. But I still find myself feeling un-confident to make deep changes and not have something subtle break on me.

I suggest setting up the following:

• It would be nice to be able to use static type-checking including axes. Seems not possible at the moment? Look into it.
• In the meantime, switch from chex to jaxtyping for more verbose array type annotations and enable jaxtyping's run-time type checking.
• Add tests.
• Code formatting with e.g. Black while I'm at it.