Type checking for pytrees

Glad to hear JAX will eventually have annotations! I was following this issue for updates. I actually started trying to add annotations myself before I realized how much work it was going to be.

One thing it would be nice to expose as soon as possible though are base types for tensors and pytrees:

Tensor = Union[np.ndarray, jnp.ndarray]  # probably needs more things in the union like tracers and DeviceArray?
PyTree = Union[Tensor,
               'PyTreeLike',
               Tuple['PyTree', ...],
               List['PyTree'],
               Dict[Hashable, 'PyTree'],
               None]

_Originally posted by @NeilGirdhar in https://github.com/google/jax/issues/1555#issuecomment-639554300_

[XuehaiPan]

Yes, in the sense that I think we’re fairly confident it cannot be done without a Python type checker that supports recursive types, which mypy and pytype (our current type checkers) do not. So, you should write Any.

mypy now supports recursive types since v0.981, and will be enabled by default since v0.990. E.g.:

JSON = Union[Dict[str, 'JSON'], List['JSON'], str, int, float, bool, None]

• python/mypy#731

[shoyer]

I can imagine two ways in which it might make sense to type check pytrees:

• Types checking over the types of leaves, e.g., a pytree of JAX arrays
• Type checking over the structure of trees, e.g., to verify that two pytrees are both dictionaries with the same keys

In theory, the first type could be checked with Generic and the second type could be checked (at least partially) with TypeVar.