Tackle Typing and Linting Errors

[gileshd]

Summary

This PR is an initial step in an overhaul of the type annotations in the code base.

Details

The aim of the PR is to:

• respond to some review comments
• fix some central core typing problems
• lay the groundwork for further typing improvements
• update typing, particularly jaxtyping, in hmm code

Further updates to other parts of the codebase will come in subsequent PRs.

Discussion

The state of PRNGKey typing in this PR is a bit inconsistent

• Sometimes keys are typed as dynamax.types.PRNGKeyT (which is an alias for jaxtyping.Array) and sometimes directly as jaxtyping.Array.
• Perhaps it would be better pick one and stick with it.
• Relevant discussion in jax docs here.

There are a lot of commmits here! After review we can perhaps squash many of the HMM typing commits into appropriate groups.

Commit Overview

This is a brief categorisation of the commits in this PR:

• Fix linting errors:
  • 3f89961 Add ruff ignore F722
    • Addresses review comment "Incorrect typing"
  • 3ad1fa7 Remove unused imports
  • dc7f675 Convert strings with latex to raw-strings
• General jaxtyping fix:
  • 20a1df7 Prepend space in uni-dim jaxtyping hints
    • Addresses review comment "Incorrect typing"
• PRNGKey typing:
  • ca9a75c Fix jr.PRNGKey type hints
  • 75205f8 Rename and change PRNGKey Type
• Dynamax typing:
  • 42c5c22 Add IntScalar type
• Misc typing:
  • 7da57ef Minor arg and type changes in utils/utils.py
• HMM Typing:
  • HMM Core:
    • 62df3ba Update HMM[Parameter|Property]Set protocols
    • 45134e8 Update type annotations in hmm base classes
    • 0d4ca74 Update type annotations in hmm inference code.
    • cfd3bfc Fix type annotations in hmm parallel inference
    • d0a0e0f Add further type annotations to hmm transitions class
    • 4b2c40b Add further type annotations to hmm initial base class
  • HMM Models:
    • 4c84ddc Add further type annotations to categorical hmm
    • 7e04c9e Add further type annotations to arhmm
    • b2fffa4 Add further type annotations to linreghmm
    • a45e61b Add further type annotations to Bernoulli HMM
    • 32db62a Add further type annotations to Gamma HMM
    • 81138f8 Add further type annotations to Gaussian HMMs
    • 4e36520 Add further type annotations to gmhmms
    • 4b7ee53 Add further type annotations to logreg hmm
    • e246606 Add further type annotations to multinomialhmm
    • e19748d Add further type annotations to poisson hmm
    • c2098cf Add further type annotations to categorical glm hmm
• Other models:
  • 547c610 Fix LinearGaussianSSM.sample type hint
  • 58c9127 Change type hints to jaxtyping in slds code