More/Better typing?

[junpenglao]

TL;dr: Should we introduce jaxtyping and beartype?

Currently Problem

As noted in #543 (comment), I am not happy about the way logdensity is annotated (e.g., see example). While they are Array (as they are in most cases should be output of a Jax function), we annotate them as float to empathizes they should be scalar.

Potential Solution

While jaxtyping is introduced to solve this, it would make the code much longer

# Accepts floating-point 2D arrays with matching dimensions
def matrix_multiply(x: Float[Array, "dim1 dim2"],
                    y: Float[Array, "dim2 dim3"]
                  ) -> Float[Array, "dim1 dim3"]:
    ...

We also have a lot of function that accepts nested structure (PyTree), although we can introduce beartype so runtime check speed would not be a concern.

[albcab]

I like runtime type-checking, it's useful for efficiently writing bug-free code and a great complement to testing. It seems like we'd need to change various float and bool to Float/Bool[Array, "1"] and we can get runtime type-checking? Or am I missing something else? If that is the case then it sounds like a good trade-off.