TypeCheckError: argument "<arg name>" (jaxlib.xla_extension.ArrayImpl) is not an instance of jax._src.prng.PRNGKeyArray about jaxtyping HOT 3 OPEN

import beartype
import jax


# Fails at static type-checking time
def f(x: jax.Array):
    pass

f(jax.random.PRNGKey(0))
# pyright says:
# Argument of type "KeyArray" cannot be assigned to parameter "x" of type "Array" in function "f"


# Fails at runtime
@beartype.beartype
def g(x: jax.random.KeyArray):
    pass

g(jax.random.PRNGKey(0))
# beartype.roar.BeartypeCallHintParamViolation: @beartyped __main__.g() parameter x="Array([0, 0], dtype=uint32)" violates type hint <class 'jax._src.prng.PRNGKeyArray'>, as <protocol "jaxlib.xla_extension.ArrayImpl"> "Array([0, 0], dtype=uint32)" not instance of <class "jax._src.prng.PRNGKeyArray">