Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Tackle Typing and Linting Errors #379

Open
wants to merge 27 commits into
base: main
Choose a base branch
from
Open

Tackle Typing and Linting Errors #379

wants to merge 27 commits into from

Commits on Oct 16, 2024

  1. Add ruff ignore F722

    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    3f89961 View commit details
    Browse the repository at this point in the history

  2. Remove unused imports

    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    3ad1fa7 View commit details
    Browse the repository at this point in the history

  3. Convert strings with latex to raw-strings 

    This stops type-checkers worrying about escape-sequences.
    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    dc7f675 View commit details
    Browse the repository at this point in the history

  4. Prepend space in uni-dim jaxtyping hints 

    Resolves F821
- see: https://docs.kidger.site/jaxtyping/faq/#flake8-or-ruff-are-throwing-an-error
    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    20a1df7 View commit details
    Browse the repository at this point in the history

  5. Fix jr.PRNGKey type hints

    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    ca9a75c View commit details
    Browse the repository at this point in the history

  6. Rename and change PRNGKey Type 

    - Rename to PRNGKeyT to differentiate from `jax.random` function
- Change type to Array
  - see https://jax.readthedocs.io/en/latest/jep/9263-typed-keys.html#type-annotations-for-prng-keys
    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    75205f8 View commit details
    Browse the repository at this point in the history

  7. Add IntScalar type

    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    42c5c22 View commit details
    Browse the repository at this point in the history

  8. Minor arg and type changes in utils/utils.py 

    - Change default arguments to floats in `monotonically_increasing` function
- Help type-checkers parse `compute_state_overlap` function
- Update `keepdims` default value in `pytree_sum` to match `jnp.sum`
    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    7da57ef View commit details
    Browse the repository at this point in the history

  9. Update HMM[Parameter|Property]Set protocols 

    - Change attributes of these protocols to read-only.
  - This change means that instances of `NamedTuple` are compatible with
    these protocols
- Add runtime_checkable decorator to Protocols
  - This allows the Protocolcs to be checked by runtime type checkers
    such as beartype which can be used to verify jaxtyping annotations
    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    62df3ba View commit details
    Browse the repository at this point in the history

  10. Update type annotations in hmm base classes

    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    45134e8 View commit details
    Browse the repository at this point in the history

  11. Update type annotations in hmm inference code. 

    Major changes:
- Fix an error in `_compute_all_transition_probs` which caused the
  return array to be too short.
  - The `filtered_probs` were being truncated twice instead of once.
- Replace `jaxtyping.Int` with `dynamax.typing.IntScalar` or `int`
  - this reflects when integer scalar arrays are accepted
  - `jaxtyping.[Dtype]` cannot be used directly for type checking
    instead they must be used as part of an array.
- Fix the shape of `transition_matrix`:
  - if transition_matrix has a leading timestep axis it should be of
    length T-1 not of length T.
- Add annotation indicating that `transition_matrix` is an optional argument
- Raise ValueError when neither `transition_matrix` or `transition_fn`
  provided.
    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    0d4ca74 View commit details
    Browse the repository at this point in the history

  12. Fix type annotations in hmm parallel inference

    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    cfd3bfc View commit details
    Browse the repository at this point in the history

  13. Add further type annotations to hmm transitions class 

    Further, enforce that either key or initial_probs specified to
initialize method.
    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    d0a0e0f View commit details
    Browse the repository at this point in the history

  14. Add further type annotations to hmm initial base class 

    Further, enforce that either key or initial_probs specified in
initialize method.
    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    4b2c40b View commit details
    Browse the repository at this point in the history

  15. Add further type annotations to categorical hmm

    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    4c84ddc View commit details
    Browse the repository at this point in the history

  16. Add further type annotations to arhmm

    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    7e04c9e View commit details
    Browse the repository at this point in the history

  17. Add further type annotations to linreghmm

    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    b2fffa4 View commit details
    Browse the repository at this point in the history

  18. Add further type annotations to Bernoulli HMM

    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    a45e61b View commit details
    Browse the repository at this point in the history

  19. Add further type annotations to Gamma HMM

    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    32db62a View commit details
    Browse the repository at this point in the history

  20. Add further type annotations to Gaussian HMMs 

    - Add annotations to the *GaussianHMMEmssions classes
- Apply some small formatting changes to reduce line length in places
    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    81138f8 View commit details
    Browse the repository at this point in the history

  21. Add further type annotations to gmhmms 

    - Add annotations to the *GaussianMixtureHMMEmssions classes
    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    4e36520 View commit details
    Browse the repository at this point in the history

  22. Add further type annotations to logreg hmm

    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    4b7ee53 View commit details
    Browse the repository at this point in the history

  23. Add further type annotations to multinomialhmm

    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    e246606 View commit details
    Browse the repository at this point in the history

  24. Add further type annotations to poisson hmm

    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    e19748d View commit details
    Browse the repository at this point in the history

  25. Add further type annotations to categorical glm hmm

    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    c2098cf View commit details
    Browse the repository at this point in the history

  26. Fix LinearGaussianSSM.sample type hint

    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    547c610 View commit details
    Browse the repository at this point in the history

  27. Change type hints to jaxtyping in slds code

    @gileshd
    gileshd committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    58c9127 View commit details
    Browse the repository at this point in the history