Pytype resolves jnp.ndarray to Any

[aslanides]

This is simple enough to reproduce:

# foo.py
import jax.numpy as jnp

def foo(x: jnp.ndarray):
  return x + 1

foo('bar')

$ pytype foo.py
# Passes type checking.

[jakevdp]

Relevant discussion here: #943

[hamzamerzic]

Hi, are there any updates on this? How difficult would it be do get the correct behavior? And in anticipation of this, is there a type that we could define that would satisfy the check? Would the following be enough?

JAXArray = Union[
    jax.interpreters.xla.DeviceArray, 
    jax.interpreters.pxla.ShardedDeviceArray, 
    jax.interpreters.batching.BatchTracer,
]

[YouJiacheng]

Now mypy can correctly check this snippet.

[hawkinsp]

There's good news and bad news on this. pytype, outside Google, does not support type checking using the types of external dependencies unless they are in the Python typeshed types. I don't think it would make sense for JAX to do that, not least because JAX changes constantly!

So if you are a downstream consumer of JAX, you cannot get pytype to type the jax API, until pytype supports external dependencies. There's nothing JAX-specific about this!

(Inside Google, pytype supports type checking of external dependencies, and indeed, in that case, this bug is already fixed.)

Since there's no action for us to take on the JAX side, closing. This will work great if and when pytype supports this.