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:
Adds the remaining dependencies that export types to additional_dependencies for mypy in the pre-commit configuration
Fixes any newly-exposed type errors
Notes on typing fixes:
There are a few cases where functions using jax were annotated as returning float, but actually return a scalar Array object. It's unfortunate that there doesn't seem to be a way to indicate that the return value is a scalar.
I tried to follow the approach outlined in https://jax.readthedocs.io/en/latest/jax.typing.html#jax-typing-best-practices to consistently annotate functions using jax (but only for cases where mypy checks were failing, for now). TL;DR inputs should typically be ArrayLike and outputs Array, and jnp.asarray used to convert inputs where necessary.
In several places we set Array = Any and link to this (now closed) JAX issue: https://github.com/google/jax/issues/943. Now that jax.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.
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-importsin the mypy invocation. See this issue comment for more details.This PR:
additional_dependenciesfor mypy in the pre-commit configurationNotes on typing fixes:
float, but actually return a scalarArrayobject. It's unfortunate that there doesn't seem to be a way to indicate that the return value is a scalar.ArrayLikeand outputsArray, andjnp.asarrayused to convert inputs where necessary.Array = Anyand link to this (now closed) JAX issue: https://github.com/google/jax/issues/943. Now thatjax.typinghas 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.