Add additional dependencies for mypy in pre-commit #1292
Merged
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Currently mypy, when run via pre-commit, treats imports from jax and other libraries as having type
Any
, resulting in it failing to catch some type errors. This is due to the combination of 1) pre-commit running checks in isolated environments without dependencies installed by default, and 2) the mypy pre-commit hook automatically passing--ignore-missing-imports
in the mypy invocation. See this issue comment for more details.This PR:
additional_dependencies
for mypy in the pre-commit configurationNotes on typing fixes:
float
, but actually return a scalarArray
object. It's unfortunate that there doesn't seem to be a way to indicate that the return value is a scalar.ArrayLike
and outputsArray
, andjnp.asarray
used to convert inputs where necessary.Array = Any
and link to this (now closed) JAX issue: One array type to rule them all! jax-ml/jax#943. Now thatjax.typing
has been introduced we could probably make all of these cases more precise, but to stay focused this PR only changes instances where mypy checks were failing.