Add type hints to model class functionality to improve code maintainability Take advantage of [jaxtyping](https://github.com/patrick-kidger/jaxtyping) for array dimension hints