[Documentation] Push `jaxtyping` over `torchtyping` for type-checking PyTorch tensors

Hi. I was trying to use beartype as a drop-in replacement for typeguard with torchtyping.

from torch import rand
from torchtyping import TensorType, patch_typeguard
from typeguard import typechecked

patch_typeguard() 

@typechecked
def func(x: TensorType["batch"],
         y: TensorType["batch"]) -> TensorType["batch"]:
    return x + y

func(rand(3), rand(1))

This raises the expected error.

TypeError: Dimension 'batch' of inconsistent size. Got both 1 and 3.

but it doesn’t work. Should we manually write patch_beartype like patch_typeguard?

from torch import rand
from torchtyping import TensorType
from beartype import beartype

@beartype
def func(x: TensorType["batch"],
         y: TensorType["batch"]) -> TensorType["batch"]:
    return x + y

func(rand(3), rand(1))

[patrick-kidger]

Update: as a quick hack to make things easier, the latest version (0.2.9) of jaxtyping should now no longer have JAX as a mandatory dependency.

If JAX is installed then the jaxtyping.{Array, PyTree} types will exist. If it isn’t, then they won’t. Either way the Float, Int etc. types will exist and should work with PyTorch tensors.

Choice of name aside, this should mean that jaxtyping is now backend-independent.

[patrick-kidger]

Yup, I’ve been thinking about how best to handle backend agnosticism. I’ve got some plans on that front, and when/if I get more time then I’d like to revisit this.