ImportError: cannot import name 'Array' from 'jaxtyping' about jaxtyping HOT 8 OPEN

Ran into this as well for pytorch. For me the solution as described on https://docs.kidger.site/jaxtyping/api/array/#array was to use torch.Tensor in place of jaxtyping.Array, like so:

from torch import Tensor
from jaxtyping import Float32

def f(x: Float32[Tensor, "dim1 dim2"]) -> Float32[Tensor, "dim1 dim2"]:
    return x

from jaxtyping.

You need to have JAX installed as well.

jaxtyping only has JAX as an optional dependency, to support also being used with PyTorch etc.

from jaxtyping.

@MilesCranmer thanks. @patrick-kidger the jax requirement was relaxed? I don't see it anymore in pyproject.toml.
If so i'll modify my code according to the syntax suggested by Miles

from jaxtyping.

Sorry, missed this question. Yes, jaxtyping no longer depends on JAX. The name is now for historical reasons only! The syntax Miles is using is correct.

from jaxtyping.

When authoring ML runtime agnostic tooling, such as a dataset, what is the correct array type to use? I cannot assume I have neither torch, jax nor tensorflow. I currently assume at least numpy and do the following, but it might not work for other use cases:

from typing import Union, TYPE_CHECKING
from jaxtyping import Float, Bool
if TYPE_CHECKING:
    from torch import Tensor
    from numpy import ndarray
    from jaxtyping import Array as JaxArray
    # TODO: tensorflow
    Array = Union[Tensor, ndarray, JaxArray]
else:
    from numpy import ndarray as Array

from jaxtyping.

Probably something like this:

from typing import Union, TYPE_CHECKING
if TYPE_CHECKING:
    from torch import TorchTensor
    from numpy import ndarray
    from jaxtyping import Array as JaxArray
    from tensorflow import TfTensor
    Array = Union[TorchTensor, ndarray, JaxArray, TfTensor]
else:
    arrays = []
    try:
        from torch import Tensor as TorchTensor
    except Exception:
        pass
    else:
        arrays.append(TorchTensor)
    try:
        from numpy import ndarray
    except Exception:
        pass
    else:
        arrays.append(ndarray)
    try:
        from jaxtyping import Array as JaxArray
    except Exception:
        pass
    else:
        arrays.append(JaxArray)
    try:
        from tensorflow import Tensor as TfTensor
    except Exception:
        pass
    else:
        arrays.append(TfTensor)
    Array = Union[tuple(arrays)]

from jaxtyping.

Neat! I'd go for except (ModuleNotFoundError, ImportError): 😉

And it doesn't exactly roll of the tongue. Any chance this could be added to jaxtyping?

from jaxtyping.

Actually, the more general Exception is deliberate. There are cases when try-importing a module can result in other issues too, c.f. https://github.com/google/jaxtyping/blob/7a84b27da9e57c425ce4e6333121c3cdf2e07302/jaxtyping/_array_types.py#L33-L39

As for adding the above to jaxtyping. jaxtyping tries to essentially be backend-agnostic. In particular, I don't think I'd want to hardcode that it'll look for specifically torch+numpy+tensorflow+jax and nothing else. As such I think something like this is out-of-scope for jaxtyping, I'm afraid.

from jaxtyping.