Typing (type hinting) in Haiku?

Hey all!

I like using mypy as a smell test for my code and I was wondering if Haiku (and to a larger extent, Jax) has (well supported) built in typing? I tried doing

import jax.np as jnp
type(jnp.ones([3]))

I found Type Hints but it seems to only be a subset of supported ops. Does the team have any recommendations?

[tomhennigan]

We tend to use aliases internally. There is a mixture of people aliasing jnp.ndarray and chex.Array.

I know that JAX team have plans to make rng keys a distinct type (PRNGKeyArray) which is a super common thing users want to type hint (currently jax.random.PRNGKey is a function not a type).

I am not aware of a good way to type a function that accepts a PyTree of arrays in a way that will actually make the type checker do something useful.

[tomhennigan]

You could look at https://github.com/deepmind/tensor_annotations if you want a more complete type annotation (e.g. with rank and axis names). The JAX team also have some thoughts on import time type checking in JAX. So far we have not adopted any of these solutions in our libraries, there is definitely some energy from users to improve the state of typing in JAX libraries, however for now there is not a canonical solution that pleases most.