pyro-ppl / numpyro / 15420977552

93%

push

github

feat(gh-299): Type hints in distributions modules (#2032)

* feat(deps): add jaxtyping for type hinting support

* Refactor distribution module with type annotations and improve validation

- Added type annotations to various methods and properties in the Distribution class and its subclasses for better type checking and clarity.
- Updated the `enable_validation` and `validation_enabled` functions to use type hints.
- Enhanced the `arg_constraints`, `support`, and other attributes with appropriate types.
- Modified the `__init__` methods of several distribution classes to include type hints for parameters.
- Improved the `log_prob`, `sample`, and other methods to specify return types.
- Refactored the `clamp_probs` function in the util module to include type annotations.

* Enhance type annotations and add typing protocols for distributions

- Updated the `LeftTruncatedDistribution`, `RightTruncatedDistribution`, and `TwoSidedTruncatedDistribution` classes to include type hints for parameters and return types.
- Introduced a new `ConstraintLike` protocol in `numpyro/typing.py` to standardize constraint types across distributions.
- Added type hints for the `TruncatedDistribution`, `TruncatedCauchy`, `TruncatedNormal`, `TruncatedPolyaGamma`, `DoublyTruncatedPowerLaw`, and `LowerTruncatedPowerLaw` classes.
- Improved type safety and clarity in the distribution methods by specifying expected types for inputs and outputs.

* refactor(typing): move protocols to _typing.py and update imports across modules

* refactor: update set_default_validate_args to accept only bool type

* refactor: add type annotations for constraint variables in constraints.py

* fix: use `Union` type instead of `|` operator for python3.9 compatible types

* chore: isolate unit tests for custom types by numpyro

* chore: remove jaxtyping and its types with appropriate jax native types

* fix(test): import jax for compatibility and update xfail condition

* refactor: rename Protocols... (continued)

996 of 996 new or added lines in 14 files covered. (100.0%)

13653 of 14702 relevant lines covered (92.86%)

1.98 hits per line