Batched type about jaxtyping HOT 2 OPEN

With a bit of trickery, this is possible!

from typing import TypeVar, Union, TYPE_CHECKING
from jaxtyping import Array, Float32, Shaped

if TYPE_CHECKING:  # needed for static type checking compatibility
    class _Unused:
        pass

    T = TypeVar("T")
    Batched = Union[T, _Unused]
else:
    class Batched:
        def __class_getitem__(cls, item):
            return Shaped[item, "N"]

This doesn't support Batched[PyTree[...]] though; you should write that as PyTree[Batched[...]]. (With a bit of hackery you could detect Batched[PyTree[...]] in __class_getitem__ and convert it, if you really wanted to.)

Hi Patrick, thanks a lot for your quick reply! This looks great. I was playing around with this since I also have some PyTree[leaf] type-aliases, so I modified the class-getitem dunder by extracting the PyTree leaf if possible.

class Batched:
    def __class_getitem__(cls, item):
        if hasattr(item, 'leaftype'):
            return PyTree[Shaped[item.leaftype, 'N']]
        
        return Shaped[item, 'N']

from jaxtyping.

With a bit of trickery, this is possible!

from typing import TypeVar, Union, TYPE_CHECKING
from jaxtyping import Array, Float32, Shaped

if TYPE_CHECKING:  # needed for static type checking compatibility
    class _Unused:
        pass

    T = TypeVar("T")
    Batched = Union[T, _Unused]
else:
    class Batched:
        def __class_getitem__(cls, item):
            return Shaped[item, "N"]

This doesn't support Batched[PyTree[...]] though; you should write that as PyTree[Batched[...]]. (With a bit of hackery you could detect Batched[PyTree[...]] in __class_getitem__ and convert it, if you really wanted to.)

from jaxtyping.