Random key does not typecheck Key[Scalar, ""]

[EtaoinWu]

Problem

Code:

import jax.random as random
from jaxtyping import Key, Scalar

k = random.key(1)
print(isinstance(k, Key[Scalar, ""]))

Output should be True. Instead it is False.

Potential reason

Commit 1e5229c added these lines:

def _is_jax_extended_dtype(dtype: Any) -> bool:
    if not has_jax:
        return False
    try:
        is_dtype = issubclass(dtype, jax.numpy.generic)
    except TypeError:
        # `dtype` not a class
        return False
    else:
        if is_dtype:
            if hasattr(jax.dtypes, "extended"):  # jax>=0.4.14
                return jax.numpy.issubdtype(dtype, jax.dtypes.extended)
            else:  # jax<=0.4.13
                return jax.core.is_opaque_dtype(dtype)
        else:
            return False

However, in this case our object's dtype is key<fry>. It is not a class, but rather an object with type jax._src.prng.KeyTy, so this function returns False.

Due to the faulty boolean value, jaxtyping chose a wrong way to stringify the dtype (should be key<fry> as in the old versions; got prng_key), and the regex match fails for Key.

Suggested Fix: Semantically it should return True since KeyTy is derived from jax._src.dtypes.ExtendedDType.

Environment

Windows x64, but platform should be irrevalent.

micromamba create -n jaxtest python=3.12 # or use conda/mamba, should be the same
micromamba activate jaxtest
pip install jax[cpu] jaxtyping # jax==0.4.19, jaxtyping==0.2.23

[patrick-kidger]

Thanks for the report! Looks like things changed a little in underlying JAX, and we were missing a few test cases to pick up on this.

This should be fixed in #143. This will be included in the next release of jaxtyping (~3 weeks).

[patrick-kidger]

Closing as fixed in the v0.2.24 release!