Is it possible to create a parameterized type annotation? about jaxtyping HOT 2 CLOSED

Thank you so much for the speedy response!

Unfortunately I'm looking for both static and runtime checking, so I'll probably have to go with something like:

import numpy as np
from jaxtyping import Array, Float 

DataShape = "batch a b c"
DataNumPy = Float[np.ndarray, DataShape]
DataJax = Float[Array, DataShape]

x: DataNumPy = ...
y: DataJax = ...

Less elegant, but at least that should cover both static and runtime checking.

I'm not sure this is worth implementing as a feature if static checking won't work, so I'm closing the issue. But if I'm wrong and others are interested in this functionality, then please re-open and (time permitting) I'm happy to help with a PR.

Thank you for all of your work on the Jax ecosystem! Your libraries (including this one) have vastly improved my developer experience. 😄

from jaxtyping.

Ah, hmm. Yeah, the approach you're trying isn't possible at the moment.

Depending what you're after, it may be possible to make this work.

1. If you want static type-checking, then I'm afraid this is a no-go. To the best of my knowledge, what you're trying to do is impossible within the limits of Python's static type system.
2. If you want runtime type-checking, then this should be doable. We'd need to add a _MetaAbstractArray.__getitem__ method that resolves type variables.

I'd be happy to accept a PR doing this. Alternatively, at least within your own code, you could use something like the following:

import numpy as np
from jaxtyping import Float

class Data:
    def __class_getitem__(cls, item):
      return Float[item, "batch a b c"]

def foo(x: Data[np.ndarray]): ...

which should work at runtime, but once again will not work with static type checkers.

from jaxtyping.