Export Array from jax without installation

[alexj-lee]

Thanks so much for this super useful package--it has been really useful for me. I am wondering if you might consider decoupling the jax.Array import from the main jax library? I realize this would be a lot of work but to the degree that I'd like to use jaxtyping for projects that may not actually use jax code in general it might increase adoption.

[patrick-kidger]

So jaxtyping should already support being used in non-JAX projects. For example:

from torch import Tensor
from jaxtyping import Float

def f(x: Float[Tensor, "foo bar"]): ...

Moreover jaxtyping does not even list JAX as a dependency.

[alexj-lee]

Agh, so sorry, I totally missed the interchangeability of Tensor and ndarray etc in the annotation. I was trying to import Array, which then gave a jax import error, under the mistaken impression that the sense in which the library is interoperable was that Array could be used to hint for both numpy, torch, etc.

Thanks so much.